Fix unit tests

Signed-off-by: Matt Lord <[email protected]>
planetscale · Jan 6, 2025 · e9319d3 · e9319d3
1 parent d3ebbe7
commit e9319d3
Show file tree

Hide file tree

Showing 21 changed files with 709 additions and 379 deletions.
diff --git a/go/test/utils/diff.go b/go/test/utils/diff.go
@@ -68,7 +68,7 @@ func MustMatchFn(ignoredFields ...string) func(t *testing.T, want, got any, errM
 		t.Helper()
 		diff := cmp.Diff(want, got, diffOpts...)
 		if diff != "" {
-			t.Fatalf("%v: (-want +got)\n%v", errMsg, diff)
+			require.FailNow(t, "%v: (-want +got)\n%v", errMsg, diff)
 		}
 	}
 }

diff --git a/go/vt/topo/test/vschema.go b/go/vt/topo/test/vschema.go
@@ -44,9 +44,12 @@ func checkVSchema(t *testing.T, ctx context.Context, ts *topo.Server) {
 		t.Error(err)
 	}
 
-	err = ts.SaveVSchema(ctx, "test_keyspace", &vschemapb.Keyspace{
-		Tables: map[string]*vschemapb.Table{
-			"unsharded": {},
+	err = ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name: "test_keyspace",
+		Keyspace: &vschemapb.Keyspace{
+			Tables: map[string]*vschemapb.Table{
+				"unsharded": {},
+			},
 		},
 	})
 	if err != nil {
@@ -64,8 +67,11 @@ func checkVSchema(t *testing.T, ctx context.Context, ts *topo.Server) {
 		t.Errorf("GetVSchema: %s, want %s", got, want)
 	}
 
-	err = ts.SaveVSchema(ctx, "test_keyspace", &vschemapb.Keyspace{
-		Sharded: true,
+	err = ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name: "test_keyspace",
+		Keyspace: &vschemapb.Keyspace{
+			Sharded: true,
+		},
 	})
 	require.NoError(t, err)
 

diff --git a/go/vt/topo/topotests/srv_vschema_test.go b/go/vt/topo/topotests/srv_vschema_test.go
@@ -21,6 +21,7 @@ import (
 
 	"google.golang.org/protobuf/proto"
 
+	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -76,7 +77,10 @@ func TestRebuildVSchema(t *testing.T) {
 	keyspace1 := &vschemapb.Keyspace{
 		Sharded: true,
 	}
-	if err := ts.SaveVSchema(ctx, "ks1", keyspace1); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks1",
+		Keyspace: keyspace1,
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.RebuildSrvVSchema(ctx, cells); err != nil {
@@ -118,7 +122,10 @@ func TestRebuildVSchema(t *testing.T) {
 			},
 		},
 	}
-	if err := ts.SaveVSchema(ctx, "ks2", keyspace2); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks2",
+		Keyspace: keyspace2,
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.RebuildSrvVSchema(ctx, []string{"cell1"}); err != nil {
@@ -182,7 +189,10 @@ func TestRebuildVSchema(t *testing.T) {
 	wanted4.RoutingRules = rr
 
 	// Delete a keyspace, checks vschema entry in map goes away.
-	if err := ts.SaveVSchema(ctx, "ks2", &vschemapb.Keyspace{}); err != nil {
+	if err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+		Name:     "ks2",
+		Keyspace: &vschemapb.Keyspace{},
+	}); err != nil {
 		t.Fatalf("SaveVSchema(ks1) failed: %v", err)
 	}
 	if err := ts.DeleteKeyspace(ctx, "ks2"); err != nil {

diff --git a/go/vt/topo/vschema.go b/go/vt/topo/vschema.go
@@ -20,8 +20,6 @@ import (
 	"context"
 	"path"
 
-	"google.golang.org/protobuf/proto"
-
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/vterrors"
 
@@ -39,11 +37,17 @@ type KeyspaceVSchemaInfo struct {
 }
 
 func (k *KeyspaceVSchemaInfo) CloneVT() *KeyspaceVSchemaInfo {
-	return &KeyspaceVSchemaInfo{
-		Name:     k.Name,
-		Keyspace: k.Keyspace.CloneVT(),
-		version:  Version(k.version),
+	if k == nil {
+		return (*KeyspaceVSchemaInfo)(nil)
+	}
+	kc := &KeyspaceVSchemaInfo{
+		Name:    k.Name,
+		version: Version(k.version),
 	}
+	if k.Keyspace != nil {
+		kc.Keyspace = k.Keyspace.CloneVT()
+	}
+	return kc
 }
 
 // SaveVSchema saves a Vschema. A valid Vschema should be passed in. It does not verify its correctness.
@@ -91,14 +95,15 @@ func (ts *Server) GetVSchema(ctx context.Context, keyspace string) (*KeyspaceVSc
 	if err != nil {
 		return nil, err
 	}
-	var vs vschemapb.Keyspace
-	err = proto.Unmarshal(data, &vs)
+
+	vs := &vschemapb.Keyspace{}
+	err = vs.UnmarshalVT(data)
 	if err != nil {
 		return nil, vterrors.Wrapf(err, "bad vschema data: %q", data)
 	}
 	return &KeyspaceVSchemaInfo{
 		Name:     keyspace,
-		Keyspace: &vs,
+		Keyspace: vs,
 		version:  version,
 	}, nil
 }

diff --git a/go/vt/vtctl/grpcvtctldserver/server.go b/go/vt/vtctl/grpcvtctldserver/server.go
@@ -937,34 +937,37 @@ func (s *VtctldServer) CreateKeyspace(ctx context.Context, req *vtctldatapb.Crea
 	}
 
 	if req.Type == topodatapb.KeyspaceType_SNAPSHOT {
-		// Copy vschema from the base keyspace.
 		bksvs, err := s.ts.GetVSchema(ctx, req.BaseKeyspace)
-		ksvs := &topo.KeyspaceVSchemaInfo{
-			Name: req.Name,
-		}
 		if err != nil {
 			log.Infof("error from GetVSchema(%v) = %v", req.BaseKeyspace, err)
 			if topo.IsErrType(err, topo.NoNode) {
 				log.Infof("base keyspace %v does not exist; continuing with bare, unsharded vschema", req.BaseKeyspace)
-				// Create an empty vschema for the keyspace.
-				ksvs.Keyspace = &vschemapb.Keyspace{
-					Sharded:  false,
-					Tables:   map[string]*vschemapb.Table{},
-					Vindexes: map[string]*vschemapb.Vindex{},
+				bksvs = &topo.KeyspaceVSchemaInfo{
+					Name: req.Name,
+					Keyspace: &vschemapb.Keyspace{
+						Sharded:  false,
+						Tables:   map[string]*vschemapb.Table{},
+						Vindexes: map[string]*vschemapb.Vindex{},
+					},
 				}
 			} else {
 				return nil, err
 			}
 		}
 
-		// Copy the vschema from the base keyspace to the new one.
-		ksvs.Keyspace = bksvs.Keyspace.CloneVT()
+		// We don't want to clone the base keyspace's key version
+		// so we do NOT call bksvs.CloneVT() here. We instead only
+		// clone the vschemapb.Keyspace field for the new snapshot
+		// keyspace.
+		sksvs := &topo.KeyspaceVSchemaInfo{
+			Name:     req.Name,
+			Keyspace: bksvs.Keyspace.CloneVT(),
+		}
 		// SNAPSHOT keyspaces are excluded from global routing.
-		ksvs.RequireExplicitRouting = true
+		sksvs.RequireExplicitRouting = true
 
-		if err = s.ts.SaveVSchema(ctx, ksvs); err != nil {
-			err = fmt.Errorf("SaveVSchema(%v) = %w", ksvs, err)
-			return nil, err
+		if err = s.ts.SaveVSchema(ctx, sksvs); err != nil {
+			return nil, fmt.Errorf("SaveVSchema(%v) = %w", sksvs, err)
 		}
 	}
 

diff --git a/go/vt/vtctl/grpcvtctldserver/server_test.go b/go/vt/vtctl/grpcvtctldserver/server_test.go
@@ -28,6 +28,8 @@ import (
 	"testing"
 	"time"
 
+	"google.golang.org/protobuf/proto"
+
 	_flag "vitess.io/vitess/go/internal/flag"
 	"vitess.io/vitess/go/vt/vtctl/reparentutil"
 	"vitess.io/vitess/go/vt/vtctl/reparentutil/policy"
@@ -608,15 +610,18 @@ func TestApplyVSchema(t *testing.T) {
 				},
 			})
 
-			origVSchema := &vschemapb.Keyspace{
-				Sharded: true,
-				Vindexes: map[string]*vschemapb.Vindex{
-					"v1": {
-						Type: "hash",
+			origVSchema := &topo.KeyspaceVSchemaInfo{
+				Name: tt.req.Keyspace,
+				Keyspace: &vschemapb.Keyspace{
+					Sharded: true,
+					Vindexes: map[string]*vschemapb.Vindex{
+						"v1": {
+							Type: "hash",
+						},
 					},
 				},
 			}
-			err := ts.SaveVSchema(ctx, tt.req.Keyspace, origVSchema)
+			err := ts.SaveVSchema(ctx, origVSchema)
 			require.NoError(t, err)
 
 			origSrvVSchema := &vschemapb.SrvVSchema{
@@ -2577,12 +2582,12 @@ func TestCreateKeyspace(t *testing.T) {
 	tests := []struct {
 		name               string
 		topo               map[string]*topodatapb.Keyspace
-		vschemas           map[string]*vschemapb.Keyspace
+		vschemas           map[string]*topo.KeyspaceVSchemaInfo
 		req                *vtctldatapb.CreateKeyspaceRequest
 		expected           *vtctldatapb.CreateKeyspaceResponse
 		shouldErr          bool
 		vschemaShouldExist bool
-		expectedVSchema    *vschemapb.Keyspace
+		expectedVSchema    *topo.KeyspaceVSchemaInfo
 	}{
 		{
 			name: "normal keyspace",
@@ -2600,8 +2605,11 @@ func TestCreateKeyspace(t *testing.T) {
 				},
 			},
 			vschemaShouldExist: true,
-			expectedVSchema: &vschemapb.Keyspace{
-				Sharded: false,
+			expectedVSchema: &topo.KeyspaceVSchemaInfo{
+				Name: "testkeyspace",
+				Keyspace: &vschemapb.Keyspace{
+					Sharded: false,
+				},
 			},
 			shouldErr: false,
 		},
@@ -2612,12 +2620,15 @@ func TestCreateKeyspace(t *testing.T) {
 					KeyspaceType: topodatapb.KeyspaceType_NORMAL,
 				},
 			},
-			vschemas: map[string]*vschemapb.Keyspace{
+			vschemas: map[string]*topo.KeyspaceVSchemaInfo{
 				"testkeyspace": {
-					Sharded: true,
-					Vindexes: map[string]*vschemapb.Vindex{
-						"h1": {
-							Type: "hash",
+					Name: "testkeyspace",
+					Keyspace: &vschemapb.Keyspace{
+						Sharded: true,
+						Vindexes: map[string]*vschemapb.Vindex{
+							"h1": {
+								Type: "hash",
+							},
 						},
 					},
 				},
@@ -2643,14 +2654,17 @@ func TestCreateKeyspace(t *testing.T) {
 				},
 			},
 			vschemaShouldExist: true,
-			expectedVSchema: &vschemapb.Keyspace{
-				Sharded: true,
-				Vindexes: map[string]*vschemapb.Vindex{
-					"h1": {
-						Type: "hash",
+			expectedVSchema: &topo.KeyspaceVSchemaInfo{
+				Name: "testkeyspace",
+				Keyspace: &vschemapb.Keyspace{
+					Sharded: true,
+					Vindexes: map[string]*vschemapb.Vindex{
+						"h1": {
+							Type: "hash",
+						},
 					},
+					RequireExplicitRouting: true,
 				},
-				RequireExplicitRouting: true,
 			},
 			shouldErr: false,
 		},
@@ -2696,9 +2710,12 @@ func TestCreateKeyspace(t *testing.T) {
 				},
 			},
 			vschemaShouldExist: true,
-			expectedVSchema: &vschemapb.Keyspace{
-				Sharded:                false,
-				RequireExplicitRouting: true,
+			expectedVSchema: &topo.KeyspaceVSchemaInfo{
+				Name: "testsnapshot",
+				Keyspace: &vschemapb.Keyspace{
+					Sharded:                false,
+					RequireExplicitRouting: true,
+				},
 			},
 			shouldErr: false,
 		},
@@ -2748,8 +2765,11 @@ func TestCreateKeyspace(t *testing.T) {
 				},
 			},
 			vschemaShouldExist: true,
-			expectedVSchema: &vschemapb.Keyspace{
-				Sharded: false,
+			expectedVSchema: &topo.KeyspaceVSchemaInfo{
+				Name: "testkeyspace",
+				Keyspace: &vschemapb.Keyspace{
+					Sharded: false,
+				},
 			},
 			shouldErr: false,
 		},
@@ -2772,7 +2792,8 @@ func TestCreateKeyspace(t *testing.T) {
 			vschemaShouldExist: false,
 			expectedVSchema:    nil,
 			shouldErr:          false,
-		}, {
+		},
+		{
 			name: "keyspace with durability policy specified",
 			topo: nil,
 			req: &vtctldatapb.CreateKeyspaceRequest{
@@ -2790,8 +2811,11 @@ func TestCreateKeyspace(t *testing.T) {
 				},
 			},
 			vschemaShouldExist: true,
-			expectedVSchema: &vschemapb.Keyspace{
-				Sharded: false,
+			expectedVSchema: &topo.KeyspaceVSchemaInfo{
+				Name: "testkeyspace",
+				Keyspace: &vschemapb.Keyspace{
+					Sharded: false,
+				},
 			},
 			shouldErr: false,
 		},
@@ -2820,7 +2844,7 @@ func TestCreateKeyspace(t *testing.T) {
 			}
 
 			for name, vs := range tt.vschemas {
-				require.NoError(t, ts.SaveVSchema(ctx, name, vs), "error in SaveVSchema(%v, %+v)", name, vs)
+				require.NoError(t, ts.SaveVSchema(ctx, vs), "error in SaveVSchema(%v, %+v)", name, vs)
 			}
 
 			// Create the keyspace and make some assertions
@@ -2829,7 +2853,6 @@ func TestCreateKeyspace(t *testing.T) {
 				assert.Error(t, err)
 				return
 			}
-
 			assert.NoError(t, err)
 			testutil.AssertKeyspacesEqual(t, tt.expected.Keyspace, resp.Keyspace, "%+v\n%+v\n", tt.expected.Keyspace, resp.Keyspace)
 
@@ -2858,7 +2881,7 @@ func TestCreateKeyspace(t *testing.T) {
 				return
 			}
 			assert.NoError(t, err)
-			utils.MustMatch(t, tt.expectedVSchema, vs)
+			require.True(t, proto.Equal(tt.expectedVSchema, vs), "expected vschema for %s: %+v, got: %+v", tt.req.Name, tt.expectedVSchema, vs)
 		})
 	}
 }
@@ -8382,11 +8405,14 @@ func TestGetVSchema(t *testing.T) {
 	})
 
 	t.Run("found", func(t *testing.T) {
-		err := ts.SaveVSchema(ctx, "testkeyspace", &vschemapb.Keyspace{
-			Sharded: true,
-			Vindexes: map[string]*vschemapb.Vindex{
-				"v1": {
-					Type: "hash",
+		err := ts.SaveVSchema(ctx, &topo.KeyspaceVSchemaInfo{
+			Name: "testkeyspace",
+			Keyspace: &vschemapb.Keyspace{
+				Sharded: true,
+				Vindexes: map[string]*vschemapb.Vindex{
+					"v1": {
+						Type: "hash",
+					},
 				},
 			},
 		})