route restartModule through viamClient dial infra and implement TestRestartModule

cli/auth.go

@@ -138,7 +138,7 @@ func (c *viamClient) loginAction(ctx context.Context, cmd *cli.Command) error {
 		already := "Already l"
 		if !alreadyLoggedIn {
 			already = "L"
-			// only print the viam logo if we are in an interative terminal
+			// only print the viam logo if we are in an interactive terminal
 			if term.IsTerminal(int(os.Stdout.Fd())) {
 				viamLogo(cmd.Root().Writer)
 			}

cli/client_test.go

@@ -193,19 +193,39 @@ func setupWithRunningPart(
 	cliArgs ...string,
 ) (*cli.Command, *viamClient, *testWriter, *testWriter) {
 	t.Helper()
+	return setupWithRunningPartAndConfig(t, asc, dataClient, buildClient, defaultFlags, authMethod, partFQDN, nil, cliArgs...)
+}
+
+// setupWithRunningPartAndConfig is like setupWithRunningPart but allows supplying a custom robot
+// config (e.g. with local modules). If robotCfg is nil, the default shell-only config is used.
+func setupWithRunningPartAndConfig(
+	t *testing.T,
+	asc apppb.AppServiceClient,
+	dataClient datapb.DataServiceClient,
+	buildClient buildpb.BuildServiceClient,
+	defaultFlags map[string]any,
+	authMethod string,
+	partFQDN string,
+	robotCfg *robotconfig.Config,
+	cliArgs ...string,
+) (*cli.Command, *viamClient, *testWriter, *testWriter) {
+	t.Helper()
 
 	cCtx, ac, out, errOut := setup(asc, dataClient, buildClient, defaultFlags, authMethod, cliArgs...)
 
-	// this config could later become a parameter
-	r, err := robotimpl.New(context.Background(), &robotconfig.Config{
-		Services: []resource.Config{
-			{
-				Name:  "shell1",
-				API:   shell.API,
-				Model: resource.DefaultServiceModel,
+	cfg := robotCfg
+	if cfg == nil {
+		cfg = &robotconfig.Config{
+			Services: []resource.Config{
+				{
+					Name:  "shell1",
+					API:   shell.API,
+					Model: resource.DefaultServiceModel,
+				},
 			},
-		},
-	}, nil, logging.NewInMemoryLogger(t))
+		}
+	}
+	r, err := robotimpl.New(context.Background(), cfg, nil, logging.NewInMemoryLogger(t))
 	test.That(t, err, test.ShouldBeNil)
 
 	options, _, addr := robottestutils.CreateBaseOptionsAndListener(t)

cli/module_build.go

@@ -27,13 +27,11 @@ import (
 	buildpb "go.viam.com/api/app/build/v1"
 	v1 "go.viam.com/api/app/packages/v1"
 	apppb "go.viam.com/api/app/v1"
-	"go.viam.com/utils/rpc"
 	"golang.org/x/exp/maps"
 
 	"go.viam.com/rdk/config"
 	"go.viam.com/rdk/logging"
 	"go.viam.com/rdk/robot"
-	"go.viam.com/rdk/robot/client"
 	"go.viam.com/rdk/utils"
 )
 
@@ -1585,24 +1583,15 @@ func restartModule(
 	if err != nil {
 		return err
 	}
-	apiRes, err := vc.client.GetRobotAPIKeys(ctx, &apppb.GetRobotAPIKeysRequest{RobotId: part.Robot})
+	args, err := getGlobalArgs(cmd)
 	if err != nil {
 		return err
 	}
-	if len(apiRes.ApiKeys) == 0 {
-		return errors.New("API keys list for this machine is empty. You can create one with \"viam machine api-key create\"")
-	}
-	key := apiRes.ApiKeys[0]
-	args, err := getGlobalArgs(cmd)
+	dialCtx, fqdn, rpcOpts, err := vc.prepareDialInner(ctx, part.Fqdn, args.Debug)
 	if err != nil {
 		return err
 	}
-	debugf(cmd.Root().Writer, args.Debug, "using API key: %s %s", key.ApiKey.Id, key.ApiKey.Name)
-	creds := rpc.WithEntityCredentials(key.ApiKey.Id, rpc.Credentials{
-		Type:    rpc.CredentialsTypeAPIKey,
-		Payload: key.ApiKey.Key,
-	})
-	robotClient, err := client.New(ctx, part.Fqdn, logger, client.WithDialOptions(creds))
+	robotClient, err := vc.connectToRobot(dialCtx, fqdn, rpcOpts, args.Debug, logger)
 	if err != nil {
 		return err
 	}

cli/module_reload_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/pkg/errors"
 	v1 "go.viam.com/api/app/build/v1"
 	apppb "go.viam.com/api/app/v1"
@@ -17,6 +18,7 @@ import (
 
 	rdkConfig "go.viam.com/rdk/config"
 	"go.viam.com/rdk/logging"
+	rtestutils "go.viam.com/rdk/testutils"
 	"go.viam.com/rdk/testutils/inject"
 )
 
@@ -309,7 +311,58 @@ func TestReloadWithCloudConfig(t *testing.T) {
 }
 
 func TestRestartModule(t *testing.T) {
-	t.Skip("restartModule test requires fake robot client")
+	logger := logging.NewTestLogger(t)
+	ctx := context.Background()
+
+	partFqdn := uuid.NewString()
+	const testPartID = "cli-restart-part"
+	simplePath := rtestutils.BuildTempModule(t, "examples/customresources/demos/simplemodule")
+	modName := "cli-restart-module"
+	robotCfg := &rdkConfig.Config{
+		Modules: []rdkConfig.Module{
+			{
+				Name:    modName,
+				ExePath: simplePath,
+				Type:    rdkConfig.ModuleTypeLocal,
+			},
+		},
+	}
+	emptyConf, err := structpb.NewStruct(map[string]any{"modules": []any{}})
+	test.That(t, err, test.ShouldBeNil)
+
+	asc := &inject.AppServiceClient{
+		GetRobotPartFunc: func(ctx context.Context, req *apppb.GetRobotPartRequest,
+			opts ...grpc.CallOption,
+		) (*apppb.GetRobotPartResponse, error) {
+			test.That(t, req.Id, test.ShouldEqual, testPartID)
+			return &apppb.GetRobotPartResponse{
+				Part: &apppb.RobotPart{
+					Id:          testPartID,
+					Fqdn:        partFqdn,
+					RobotConfig: emptyConf,
+					LastUpdated: timestamppb.New(time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)),
+				},
+				ConfigJson: ``,
+			}, nil
+		},
+	}
+
+	flags := map[string]any{
+		generalFlagPartID: testPartID,
+		generalFlagName:   modName,
+	}
+
+	cCtx, vc, _, _ := setupWithRunningPartAndConfig(
+		t, asc, nil, &inject.BuildServiceClient{},
+		flags, "token", partFqdn, robotCfg,
+	)
+	test.That(t, vc.loginAction(ctx, cCtx), test.ShouldBeNil)
+
+	partResp, err := vc.getRobotPart(ctx, testPartID)
+	test.That(t, err, test.ShouldBeNil)
+
+	err = restartModule(ctx, cCtx, vc, partResp.Part, nil, logger)
+	test.That(t, err, test.ShouldBeNil)
 }
 
 func TestResolvePartId(t *testing.T) {