diff --git a/go/vt/vtctl/endtoend/get_schema_test.go b/go/vt/vtctl/endtoend/get_schema_test.go index ddb0c204048..1e93ce633c3 100644 --- a/go/vt/vtctl/endtoend/get_schema_test.go +++ b/go/vt/vtctl/endtoend/get_schema_test.go @@ -2,7 +2,6 @@ package endtoend import ( "context" - "fmt" "testing" "github.com/google/uuid" @@ -14,7 +13,7 @@ import ( "vitess.io/vitess/go/vt/topo/memorytopo" "vitess.io/vitess/go/vt/topo/topoproto" "vitess.io/vitess/go/vt/vtctl" - "vitess.io/vitess/go/vt/vttablet/faketmclient" + "vitess.io/vitess/go/vt/vtctl/grpcvtctldserver/testutil" "vitess.io/vitess/go/vt/vttablet/tmclient" "vitess.io/vitess/go/vt/wrangler" @@ -23,29 +22,6 @@ import ( topodatapb "vitess.io/vitess/go/vt/proto/topodata" ) -type fakeTabletManagerClient struct { - tmclient.TabletManagerClient - schemas map[string]*tabletmanagerdatapb.SchemaDefinition -} - -func newTMClient() *fakeTabletManagerClient { - return &fakeTabletManagerClient{ - TabletManagerClient: faketmclient.NewFakeTabletManagerClient(), - schemas: map[string]*tabletmanagerdatapb.SchemaDefinition{}, - } -} - -func (c *fakeTabletManagerClient) GetSchema(ctx context.Context, tablet *topodatapb.Tablet, tablets []string, excludeTables []string, includeViews bool) (*tabletmanagerdatapb.SchemaDefinition, error) { - key := topoproto.TabletAliasString(tablet.Alias) - - schema, ok := c.schemas[key] - if !ok { - return nil, fmt.Errorf("no schemas for %s", key) - } - - return schema, nil -} - func TestGetSchema(t *testing.T) { ctx := context.Background() @@ -162,8 +138,10 @@ func TestGetSchema(t *testing.T) { }, } - tmc := newTMClient() - tmc.schemas[topoproto.TabletAliasString(tablet.Alias)] = sd + tmc := testutil.TabletManagerClient + tmc.Schemas[topoproto.TabletAliasString(tablet.Alias)] = sd + + *tmclient.TabletManagerProtocol = testutil.TabletManagerClientProtocol logger := logutil.NewMemoryLogger()