Simple_Schema_Changes: WaitGroup

flow/e2e/bigquery/peer_flow_bq_test.go

@@ -839,6 +839,9 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Simple_Schema_Changes_BQ() {
 		MaxBatchSize:     100,
 	}
 
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
 	// in a separate goroutine, wait for PeerFlowStatusQuery to finish setup
 	// and then insert and mutate schema repeatedly.
 	go func() {
@@ -894,6 +897,8 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Simple_Schema_Changes_BQ() {
 		// verify we got our two rows, if schema did not match up it will error.
 		e2e.NormalizeFlowCountQuery(env, connectionGen, 8)
 		s.compareTableContentsBQ("test_simple_schema_changes", "id,c1")
+
+		wg.Done()
 	}()
 
 	env.ExecuteWorkflow(peerflow.CDCFlowWorkflowWithConfig, flowConnConfig, &limits, nil)
@@ -905,6 +910,7 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Simple_Schema_Changes_BQ() {
 	// allow only continue as new error
 	require.Contains(s.t, err.Error(), "continue as new")
 
+	wg.Wait()
 	env.AssertExpectations(s.t)
 }
 

flow/e2e/snowflake/peer_flow_sf_test.go

@@ -845,6 +845,9 @@ func (s PeerFlowE2ETestSuiteSF) Test_Simple_Schema_Changes_SF() {
 		MaxBatchSize:     100,
 	}
 
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
 	// in a separate goroutine, wait for PeerFlowStatusQuery to finish setup
 	// and then insert and mutate schema repeatedly.
 	go func() {
@@ -990,6 +993,8 @@ func (s PeerFlowE2ETestSuiteSF) Test_Simple_Schema_Changes_SF() {
 		require.NoError(s.t, err)
 		require.Equal(s.t, expectedTableSchema, output.TableNameSchemaMapping[dstTableName])
 		s.compareTableContentsSF("test_simple_schema_changes", "id,c1")
+
+		wg.Done()
 	}()
 
 	env.ExecuteWorkflow(peerflow.CDCFlowWorkflowWithConfig, flowConnConfig, &limits, nil)
@@ -1001,6 +1006,7 @@ func (s PeerFlowE2ETestSuiteSF) Test_Simple_Schema_Changes_SF() {
 	// allow only continue as new error
 	require.Contains(s.t, err.Error(), "continue as new")
 
+	wg.Wait()
 	env.AssertExpectations(s.t)
 }