Fix unit test

Signed-off-by: Matt Lord <[email protected]>
vitessio · Jan 2, 2025 · 8d8b3b4 · 8d8b3b4
1 parent 0d136d4
commit 8d8b3b4
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 2 deletions.
diff --git a/go/cmd/vtctldclient/command/root_test.go b/go/cmd/vtctldclient/command/root_test.go
@@ -30,6 +30,7 @@ import (
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 	"vitess.io/vitess/go/vt/vtctl/localvtctldclient"
+	"vitess.io/vitess/go/vt/vttablet/tmclient"
 
 	vtctlservicepb "vitess.io/vitess/go/vt/proto/vtctlservice"
 )
@@ -67,15 +68,18 @@ func TestRootWithInternalVtctld(t *testing.T) {
 	cell := "zone1"
 	ts, factory := memorytopo.NewServerAndFactory(ctx, cell)
 	topo.RegisterFactory("test", factory)
+	origProtocol := command.VtctldClientProtocol
 	command.VtctldClientProtocol = "local"
+	tmclient.RegisterTabletManagerClientFactory("grpc", func() tmclient.TabletManagerClient {
+		return nil
+	})
 	baseArgs := []string{"vtctldclient", "--server", "internal", "--topo-implementation", "test"}
 
 	args := append([]string{}, os.Args...)
-	protocol := command.VtctldClientProtocol
 	t.Cleanup(func() {
 		ts.Close()
 		os.Args = append([]string{}, args...)
-		command.VtctldClientProtocol = protocol
+		command.VtctldClientProtocol = origProtocol
 	})
 
 	testCases := []struct {

diff --git a/go/cmd/vtctldclient/main.go b/go/cmd/vtctldclient/main.go
@@ -30,6 +30,7 @@ import (
 	"vitess.io/vitess/go/vt/vtctl/grpcclientcommon"
 	"vitess.io/vitess/go/vt/vtctl/vtctlclient"
 	"vitess.io/vitess/go/vt/vttablet/grpctmclient"
+	"vitess.io/vitess/go/vt/vttablet/tmclient"
 
 	_flag "vitess.io/vitess/go/internal/flag"
 )
@@ -46,6 +47,7 @@ func main() {
 	grpccommon.RegisterFlags(command.Root.PersistentFlags())
 	grpcclientcommon.RegisterFlags(command.Root.PersistentFlags())
 	grpctmclient.RegisterFlags(command.Root.PersistentFlags())
+	tmclient.RegisterFlags(command.Root.PersistentFlags())
 	servenv.RegisterMySQLServerFlags(command.Root.PersistentFlags())
 	vtctlclient.RegisterFlags(command.Root.PersistentFlags())
 	acl.RegisterFlags(command.Root.PersistentFlags())

diff --git a/go/cmd/vtctldclient/plugin_grpctmclient.go b/go/cmd/vtctldclient/plugin_grpctmclient.go
@@ -20,4 +20,5 @@ package main
 
 import (
 	_ "vitess.io/vitess/go/vt/vttablet/grpctmclient"
+	_ "vitess.io/vitess/go/vt/vttablet/tmclient"
 )
diff --git a/go/flags/endtoend/vtctldclient.txt b/go/flags/endtoend/vtctldclient.txt
@@ -150,6 +150,7 @@ Flags:
       --tablet_manager_grpc_crl string           the server crl to use to validate server certificates when connecting
       --tablet_manager_grpc_key string           the key to use to connect
       --tablet_manager_grpc_server_name string   the server name to use to validate server certificate
+      --tablet_manager_protocol string           Protocol to use to make tabletmanager RPCs to vttablets. (default "grpc")
       --topo-global-root string                  the path of the global topology data in the global topology server (default "/vitess/global")
       --topo-global-server-address strings       the address of the global topology server(s) (default [localhost:2379])
       --topo-implementation string               the topology implementation to use (default "etcd2")