feat: Enhance test suite (#1523)

Enhance the destination test suite
- Run a double migration when syncing a safe migration. This will catch any issues or incomplete migrations that occurred
- Add more tests cases for the automigration to `_cq_id` 

In order to enable this functionality the `TestDataGenerator` needed to be enhanced to enable setting of the `seed` and `counter` so that records would not collide when tests use multiple passes

Author: bbernays

plugin/nulls_test.go

@@ -14,7 +14,7 @@ import (
 func TestWithTestIgnoreNullsInLists(t *testing.T) {
 	s := &WriterTestSuite{ignoreNullsInLists: true}
 
-	tg := schema.NewTestDataGenerator()
+	tg := schema.NewTestDataGenerator(0)
 	source := schema.TestTable("ignore_nulls_in_lists", schema.TestSourceOptions{})
 	resource := s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
 		SourceName: "allow_null",
@@ -59,7 +59,7 @@ func TestWithTestSourceAllowNull(t *testing.T) {
 		}
 	}}
 
-	tg := schema.NewTestDataGenerator()
+	tg := schema.NewTestDataGenerator(0)
 	source := schema.TestTable("allow_null", schema.TestSourceOptions{})
 	resource := s.handleNulls(tg.Generate(source, schema.GenTestDataOptions{
 		SourceName: "allow_null",

plugin/testing_write_delete.go

@@ -91,7 +91,7 @@ func (s *WriterTestSuite) testDeleteStaleAll(ctx context.Context, t *testing.T)
 	table.Columns[table.Columns.Index("id")].PrimaryKey = true
 	r.NoErrorf(s.plugin.writeOne(ctx, &message.WriteMigrateTable{Table: table}), "failed to create table")
 
-	tg := schema.NewTestDataGenerator()
+	tg := schema.NewTestDataGenerator(0)
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:       rowsPerRecord,
 		TimePrecision: s.genDatOptions.TimePrecision,

plugin/testing_write_insert.go

@@ -89,7 +89,7 @@ func (s *WriterTestSuite) testInsertAll(ctx context.Context) error {
 	}); err != nil {
 		return fmt.Errorf("failed to create table: %w", err)
 	}
-	tg := schema.NewTestDataGenerator()
+	tg := schema.NewTestDataGenerator(0)
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:       rowsPerRecord,
 		TimePrecision: s.genDatOptions.TimePrecision,

plugin/testing_write_migrate.go

@@ -21,14 +21,20 @@ func tableUUIDSuffix() string {
 
 // nolint:revive
 func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, source *schema.Table, supportsSafeMigrate bool, writeOptionMigrateForce bool) error {
-	const rowsPerRecord = 10
+	var rowsPerRecord = int(10)
 	if err := s.plugin.writeOne(ctx, &message.WriteMigrateTable{
 		Table:        source,
 		MigrateForce: writeOptionMigrateForce,
 	}); err != nil {
 		return fmt.Errorf("failed to create table: %w", err)
 	}
 
+	records, err := s.plugin.readAll(ctx, source)
+	if err != nil {
+		return fmt.Errorf("failed to read initial records: %w", err)
+	}
+	initialItems := int(TotalRows(records))
+
 	sourceName := target.Name
 	syncTime := time.Now().UTC().Round(1 * time.Second)
 	opts := schema.GenTestDataOptions{
@@ -37,7 +43,10 @@ func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, sou
 		MaxRows:       rowsPerRecord,
 		TimePrecision: s.genDatOptions.TimePrecision,
 	}
-	tg := schema.NewTestDataGenerator()
+	// Test Generator should be initialized with the current number of items in the destination
+	// this allows us to have multi-pass tests that ensure the migrations are stable
+	// create--> write --> migrate --> write -->migrate -->write-->migrate -->write
+	tg := schema.NewTestDataGenerator(uint64(initialItems))
 	resource1 := tg.Generate(source, opts)
 	if err := s.plugin.writeOne(ctx, &message.WriteInsert{
 		Record: resource1,
@@ -46,15 +55,18 @@ func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, sou
 	}
 	resource1 = s.handleNulls(resource1) // we process nulls after writing
 
-	records, err := s.plugin.readAll(ctx, source)
+	records, err = s.plugin.readAll(ctx, source)
 	if err != nil {
 		return fmt.Errorf("failed to sync: %w", err)
 	}
+	sortRecords(source, records, "id")
+	records = records[initialItems:]
+
 	totalItems := TotalRows(records)
-	if totalItems != rowsPerRecord {
+	if totalItems != int64(rowsPerRecord) {
 		return fmt.Errorf("expected items: %d, got: %d", rowsPerRecord, totalItems)
 	}
-	sortRecords(source, records, "id")
+
 	if diff := RecordsDiff(source.ToArrowSchema(), records, []arrow.Record{resource1}); diff != "" {
 		return fmt.Errorf("first record differs from expectation: %s", diff)
 	}
@@ -79,14 +91,14 @@ func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, sou
 		return fmt.Errorf("failed to readAll: %w", err)
 	}
 	sortRecords(target, records, "id")
-
+	records = records[initialItems:]
 	lastRow := resource2.NewSlice(resource2.NumRows()-1, resource2.NumRows())
 	// if force migration is not required, we don't expect any items to be dropped (so there should be 2 items)
 	if !writeOptionMigrateForce || supportsSafeMigrate {
-		if err := expectRows(target.ToArrowSchema(), records, 2*rowsPerRecord, lastRow); err != nil {
-			if writeOptionMigrateForce && TotalRows(records) == rowsPerRecord {
+		if err := expectRows(target.ToArrowSchema(), records, 2*int64(rowsPerRecord), lastRow); err != nil {
+			if writeOptionMigrateForce && TotalRows(records) == int64(rowsPerRecord) {
 				// if force migration is required, we can also expect 1 item to be dropped
-				return expectRows(target.ToArrowSchema(), records, rowsPerRecord, lastRow)
+				return expectRows(target.ToArrowSchema(), records, int64(rowsPerRecord), lastRow)
 			}
 
 			return err
@@ -95,7 +107,7 @@ func (s *WriterTestSuite) migrate(ctx context.Context, target *schema.Table, sou
 		return nil
 	}
 
-	return expectRows(target.ToArrowSchema(), records, rowsPerRecord, lastRow)
+	return expectRows(target.ToArrowSchema(), records, int64(rowsPerRecord), lastRow)
 }
 
 // nolint:revive
@@ -112,7 +124,7 @@ func (s *WriterTestSuite) testMigrate(
 		if !forceMigrate && !s.tests.SafeMigrations.AddColumn {
 			t.Skip("skipping test: add_column")
 		}
-		tableName := "add_column" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_add_column" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -129,16 +141,17 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean},
 			},
 		}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.AddColumn, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate %s: %v", tableName, err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.AddColumn, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
 	})
 
 	t.Run("add_column_not_null"+suffix, func(t *testing.T) {
 		if !forceMigrate && !s.tests.SafeMigrations.AddColumnNotNull {
 			t.Skip("skipping test: add_column_not_null")
 		}
-		tableName := "add_column_not_null" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_add_column_not_null" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -154,16 +167,18 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "uuid", Type: types.ExtensionTypes.UUID},
 				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean, NotNull: true},
 			}}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.AddColumnNotNull, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate add_column_not_null: %v", err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.AddColumnNotNull, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
+
 	})
 
 	t.Run("remove_column"+suffix, func(t *testing.T) {
 		if !forceMigrate && !s.tests.SafeMigrations.RemoveColumn {
 			t.Skip("skipping test: remove_column")
 		}
-		tableName := "remove_column" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_remove_column" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -177,16 +192,17 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "id", Type: arrow.PrimitiveTypes.Int64},
 				{Name: "uuid", Type: types.ExtensionTypes.UUID},
 			}}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.RemoveColumn, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate remove_column: %v", err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.RemoveColumn, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
 	})
 
 	t.Run("remove_column_not_null"+suffix, func(t *testing.T) {
 		if !forceMigrate && !s.tests.SafeMigrations.RemoveColumnNotNull {
 			t.Skip("skipping test: remove_column_not_null")
 		}
-		tableName := "remove_column_not_null" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_remove_column_not_null" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -201,16 +217,17 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "id", Type: arrow.PrimitiveTypes.Int64},
 				{Name: "uuid", Type: types.ExtensionTypes.UUID},
 			}}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.RemoveColumnNotNull, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate remove_column_not_null: %v", err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.RemoveColumnNotNull, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
 	})
 
 	t.Run("change_column"+suffix, func(t *testing.T) {
 		if !forceMigrate && !s.tests.SafeMigrations.ChangeColumn {
 			t.Skip("skipping test: change_column")
 		}
-		tableName := "change_column" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_change_column" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -225,16 +242,17 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "uuid", Type: types.ExtensionTypes.UUID},
 				{Name: "bool", Type: arrow.BinaryTypes.String, NotNull: true},
 			}}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.ChangeColumn, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate change_column: %v", err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.ChangeColumn, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
 	})
 
 	t.Run("move_to_cq_id_only"+suffix, func(t *testing.T) {
 		if !forceMigrate && !s.tests.SafeMigrations.MovePKToCQOnly {
 			t.Skip("skipping test: move_to_cq_id_only")
 		}
-		tableName := "move_to_cq_id_only" + suffix + "_" + tableUUIDSuffix()
+		tableName := "cq_move_to_cq_id_only" + suffix + "_" + tableUUIDSuffix()
 		source := &schema.Table{
 			Name: tableName,
 			Columns: schema.ColumnList{
@@ -251,16 +269,43 @@ func (s *WriterTestSuite) testMigrate(
 				{Name: "uuid", Type: types.ExtensionTypes.UUID},
 				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean, NotNull: true},
 			}}
-		if err := s.migrate(ctx, target, source, s.tests.SafeMigrations.MovePKToCQOnly, forceMigrate); err != nil {
-			t.Fatalf("failed to migrate move_to_cq_id_only: %v", err)
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.MovePKToCQOnly, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
+		}
+	})
+	t.Run("move_to_cq_id_only_adding_pkc"+suffix, func(t *testing.T) {
+		if !forceMigrate && !s.tests.SafeMigrations.MovePKToCQOnly {
+			t.Skip("skipping test: move_to_cq_id_only_adding_pk")
+		}
+		tableName := "cq_move_to_cq_id_only_adding_pkc" + suffix + "_" + tableUUIDSuffix()
+		source := &schema.Table{
+			Name: tableName,
+			Columns: schema.ColumnList{
+				{Name: "_cq_id", Type: types.ExtensionTypes.UUID, NotNull: true, Unique: true},
+				{Name: "id", Type: arrow.PrimitiveTypes.Int64, PrimaryKey: true},
+				{Name: "uuid", Type: types.ExtensionTypes.UUID},
+				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean, NotNull: true, PrimaryKeyComponent: true},
+			}}
+		target := &schema.Table{
+			Name: tableName,
+			Columns: schema.ColumnList{
+				{Name: "_cq_id", Type: types.ExtensionTypes.UUID, NotNull: true, Unique: true, PrimaryKey: true},
+				{Name: "id", Type: arrow.PrimitiveTypes.Int64},
+				{Name: "uuid", Type: types.ExtensionTypes.UUID},
+				{Name: "bool", Type: arrow.FixedWidthTypes.Boolean, NotNull: true},
+			}}
+		require.NoError(t, s.migrate(ctx, target, source, s.tests.SafeMigrations.MovePKToCQOnly, forceMigrate))
+		if !forceMigrate {
+			require.NoError(t, s.migrate(ctx, target, target, true, false))
 		}
 	})
 
 	t.Run("double_migration", func(t *testing.T) {
 		if forceMigrate {
 			t.Skip("double migration test has sense only for safe migrations")
 		}
-		tableName := "double_migration_" + tableUUIDSuffix()
+		tableName := "cq_double_migration_" + tableUUIDSuffix()
 		table := schema.TestTable(tableName, s.genDatOptions)
 		// s.migrate will perform create->write->migrate->write
 		require.NoError(t, s.migrate(ctx, table, table, true, false))

plugin/testing_write_upsert.go

@@ -78,7 +78,7 @@ func (s *WriterTestSuite) testUpsertAll(ctx context.Context) error {
 		return fmt.Errorf("failed to create table: %w", err)
 	}
 
-	tg := schema.NewTestDataGenerator()
+	tg := schema.NewTestDataGenerator(0)
 	normalRecord := tg.Generate(table, schema.GenTestDataOptions{
 		MaxRows:       rowsPerRecord,
 		TimePrecision: s.genDatOptions.TimePrecision,

schema/testdata.go

@@ -204,11 +204,10 @@ type TestDataGenerator struct {
 	colToRnd map[string]*rand.Rand
 }
 
-func NewTestDataGenerator() *TestDataGenerator {
-	const seed = 0 // we can add an argument for this in the future, if necessary
+func NewTestDataGenerator(randomSeed uint64) *TestDataGenerator {
 	return &TestDataGenerator{
-		counter:  0,
-		seed:     seed,
+		counter:  int(randomSeed),
+		seed:     randomSeed,
 		colToRnd: map[string]*rand.Rand{},
 	}
 }

schema/testdata_test.go

@@ -49,6 +49,6 @@ func checkColumnsDontExist(t *testing.T, list ColumnList, cols []string) {
 func TestGenTestData(*testing.T) {
 	table := TestTable("test", TestSourceOptions{})
 	// smoke test that no panics
-	tg := NewTestDataGenerator()
+	tg := NewTestDataGenerator(0)
 	_ = tg.Generate(table, GenTestDataOptions{})
 }

schema/validators_test.go

@@ -13,7 +13,7 @@ import (
 
 func TestFindEmptyColumns(t *testing.T) {
 	table := TestTable("test", TestSourceOptions{})
-	tg := NewTestDataGenerator()
+	tg := NewTestDataGenerator(0)
 	record := tg.Generate(table, GenTestDataOptions{
 		MaxRows:  1,
 		NullRows: true,
@@ -25,7 +25,7 @@ func TestFindEmptyColumns(t *testing.T) {
 
 func TestFindEmptyColumnsNotEmpty(t *testing.T) {
 	table := TestTable("test", TestSourceOptions{})
-	tg := NewTestDataGenerator()
+	tg := NewTestDataGenerator(0)
 	record := tg.Generate(table, GenTestDataOptions{
 		MaxRows:  1,
 		NullRows: false,