fix merge tests

PeerDB-io · Jan 6, 2024 · 89af8f9 · 89af8f9
1 parent c67e5e1
commit 89af8f9
Showing 1 changed file with 30 additions and 18 deletions.
diff --git a/flow/connectors/bigquery/merge_stmt_generator_test.go b/flow/connectors/bigquery/merge_stmt_generator_test.go
@@ -9,36 +9,42 @@ import (
 )
 
 func TestGenerateUpdateStatement_WithUnchangedToastCols(t *testing.T) {
-	m := &mergeStmtGenerator{}
+	m := &mergeStmtGenerator{
+		shortColumn: map[string]string{
+			"col1": "_c0",
+			"col2": "_c1",
+			"col3": "_c2",
+		},
+	}
 	allCols := []string{"col1", "col2", "col3"}
 	unchangedToastCols := []string{"", "col2, col3", "col2", "col3"}
 
 	expected := []string{
 		"WHEN MATCHED AND _rt!=2 AND _ut=''" +
-			" THEN UPDATE SET `col1`=_d.col1,`col2`=_d.col2,`col3`=_d.col3," +
+			" THEN UPDATE SET `col1`=_d._c0,`col2`=_d._c1,`col3`=_d._c2," +
 			"`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE",
 		"WHEN MATCHED AND _rt=2 " +
 			"AND _ut='' " +
-			"THEN UPDATE SET `col1`=_d.col1,`col2`=_d.col2," +
-			"`col3`=_d.col3,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
+			"THEN UPDATE SET `col1`=_d._c0,`col2`=_d._c1," +
+			"`col3`=_d._c2,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
 		"WHEN MATCHED AND _rt!=2 AND _ut='col2,col3' " +
-			"THEN UPDATE SET `col1`=_d.col1,`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE ",
+			"THEN UPDATE SET `col1`=_d._c0,`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE ",
 		"WHEN MATCHED AND _rt=2 AND _ut='col2,col3' " +
-			"THEN UPDATE SET `col1`=_d.col1,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
+			"THEN UPDATE SET `col1`=_d._c0,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
 		"WHEN MATCHED AND _rt!=2 " +
 			"AND _ut='col2' " +
-			"THEN UPDATE SET `col1`=_d.col1,`col3`=_d.col3," +
+			"THEN UPDATE SET `col1`=_d._c0,`col3`=_d._c2," +
 			"`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE",
 		"WHEN MATCHED AND _rt=2 " +
 			"AND _ut='col2' " +
-			"THEN UPDATE SET `col1`=_d.col1,`col3`=_d.col3," +
+			"THEN UPDATE SET `col1`=_d._c0,`col3`=_d._c2," +
 			"`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE ",
 		"WHEN MATCHED AND _rt!=2 AND _ut='col3' " +
-			"THEN UPDATE SET `col1`=_d.col1," +
-			"`col2`=_d.col2,`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE ",
+			"THEN UPDATE SET `col1`=_d._c0," +
+			"`col2`=_d._c1,`synced_at`=CURRENT_TIMESTAMP,`deleted`=FALSE ",
 		"WHEN MATCHED AND _rt=2 AND _ut='col3' " +
-			"THEN UPDATE SET `col1`=_d.col1," +
-			"`col2`=_d.col2,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
+			"THEN UPDATE SET `col1`=_d._c0," +
+			"`col2`=_d._c1,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
 	}
 
 	result := m.generateUpdateStatements(allCols, unchangedToastCols, &protos.PeerDBColumns{
@@ -58,23 +64,29 @@ func TestGenerateUpdateStatement_WithUnchangedToastCols(t *testing.T) {
 }
 
 func TestGenerateUpdateStatement_NoUnchangedToastCols(t *testing.T) {
-	m := &mergeStmtGenerator{}
+	m := &mergeStmtGenerator{
+		shortColumn: map[string]string{
+			"col1": "_c0",
+			"col2": "_c1",
+			"col3": "_c2",
+		},
+	}
 	allCols := []string{"col1", "col2", "col3"}
 	unchangedToastCols := []string{""}
 
 	expected := []string{
 		"WHEN MATCHED AND _rt!=2 " +
 			"AND _ut=''" +
 			"THEN UPDATE SET " +
-			"`col1`=_d.col1," +
-			"`col2`=_d.col2," +
-			"`col3`=_d.col3," +
+			"`col1`=_d._c0," +
+			"`col2`=_d._c1," +
+			"`col3`=_d._c2," +
 			"`synced_at`=CURRENT_TIMESTAMP," +
 			"`deleted`=FALSE",
 		"WHEN MATCHED AND" +
 			"_rt=2 AND _ut=''" +
-			"THEN UPDATE SET `col1`=_d.col1,`col2`=_d.col2, " +
-			"`col3`=_d.col3,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
+			"THEN UPDATE SET `col1`=_d._c0,`col2`=_d._c1, " +
+			"`col3`=_d._c2,`synced_at`=CURRENT_TIMESTAMP,`deleted`=TRUE",
 	}
 
 	result := m.generateUpdateStatements(allCols, unchangedToastCols,