From f804e63c447fd5a2ce413ed0c3e57ba36aa67749 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <philip@peerdb.io>
Date: Sat, 16 Mar 2024 18:10:36 +0000
Subject: [PATCH] update tests, numeric NaN/inf/-inf becomes null, not NaN

---
 flow/e2e/bigquery/bigquery_helper.go   | 23 ++++++-----------------
 flow/e2e/bigquery/peer_flow_bq_test.go | 18 +++++++++++++-----
 flow/e2e/test_utils.go                 | 17 ++++-------------
 3 files changed, 23 insertions(+), 35 deletions(-)

diff --git a/flow/e2e/bigquery/bigquery_helper.go b/flow/e2e/bigquery/bigquery_helper.go
index cf2a66fac3..7c8ab5257c 100644
--- a/flow/e2e/bigquery/bigquery_helper.go
+++ b/flow/e2e/bigquery/bigquery_helper.go
@@ -422,37 +422,26 @@ func (b *BigQueryTestHelper) CheckNull(tableName string, colName []string) (bool
 }
 
 // check if NaN, Inf double values are null
-func (b *BigQueryTestHelper) CheckDoubleValues(tableName string, c1 string, c2 string) (bool, error) {
-	command := fmt.Sprintf("SELECT %s, %s FROM `%s.%s`",
-		c1, c2, b.Config.DatasetId, tableName)
+func (b *BigQueryTestHelper) SelectRow(tableName string, cols ...string) ([]bigquery.Value, error) {
+	command := fmt.Sprintf("SELECT %s FROM `%s.%s`",
+		strings.Join(cols, ","), b.Config.DatasetId, tableName)
 	q := b.client.Query(command)
 	q.DisableQueryCache = true
 	it, err := q.Read(context.Background())
 	if err != nil {
-		return false, fmt.Errorf("failed to run command: %w", err)
+		return nil, fmt.Errorf("failed to run command: %w", err)
 	}
 
 	var row []bigquery.Value
 	for {
 		err := it.Next(&row)
 		if err == iterator.Done {
-			break
+			return row, nil
 		}
 		if err != nil {
-			return false, fmt.Errorf("failed to iterate over query results: %w", err)
+			return nil, fmt.Errorf("failed to iterate over query results: %w", err)
 		}
 	}
-
-	if len(row) == 0 {
-		return false, nil
-	}
-
-	floatArr, _ := row[1].([]float64)
-	if row[0] != nil || len(floatArr) > 0 {
-		return false, nil
-	}
-
-	return true, nil
 }
 
 func qValueKindToBqColTypeString(val qvalue.QValueKind) (string, error) {
diff --git a/flow/e2e/bigquery/peer_flow_bq_test.go b/flow/e2e/bigquery/peer_flow_bq_test.go
index 3dff6e310e..223cf0ad25 100644
--- a/flow/e2e/bigquery/peer_flow_bq_test.go
+++ b/flow/e2e/bigquery/peer_flow_bq_test.go
@@ -478,7 +478,7 @@ func (s PeerFlowE2ETestSuiteBQ) Test_NaN_Doubles_BQ() {
 	srcTableName := s.attachSchemaSuffix("test_nans_bq")
 	dstTableName := "test_nans_bq"
 	_, err := s.Conn().Exec(context.Background(), fmt.Sprintf(`
-	CREATE TABLE IF NOT EXISTS %s (id serial PRIMARY KEY,c1 double precision,c2 double precision[]);
+	CREATE TABLE IF NOT EXISTS %s (id serial PRIMARY KEY,c1 double precision,c2 double precision[],c3 numeric);
 	`, srcTableName))
 	require.NoError(s.t, err)
 
@@ -499,13 +499,21 @@ func (s PeerFlowE2ETestSuiteBQ) Test_NaN_Doubles_BQ() {
 
 	// test inserting various types
 	_, err = s.Conn().Exec(context.Background(), fmt.Sprintf(`
-		INSERT INTO %s SELECT 2, 'NaN'::double precision, '{NaN, Infinity, -Infinity}';
-		`, srcTableName))
+	INSERT INTO %s SELECT 2, 'NaN'::double precision, '{NaN, Infinity, -Infinity}', 'NaN'::numeric;
+	`, srcTableName))
 	e2e.EnvNoError(s.t, env, err)
 
 	e2e.EnvWaitFor(s.t, env, 2*time.Minute, "normalize weird floats", func() bool {
-		good, err := s.bqHelper.CheckDoubleValues(dstTableName, "c1", "c2")
-		return err == nil && good
+		row, err := s.bqHelper.SelectRow(dstTableName, "c1", "c2", "c3")
+		if err != nil {
+			return false
+		}
+		if len(row) == 0 {
+			return false
+		}
+
+		floatArr, ok := row[1].([]float64)
+		return ok && row[0] == nil && len(floatArr) == 0 && row[2] == nil
 	})
 
 	env.Cancel()
diff --git a/flow/e2e/test_utils.go b/flow/e2e/test_utils.go
index dcaef74291..d2bcef6acb 100644
--- a/flow/e2e/test_utils.go
+++ b/flow/e2e/test_utils.go
@@ -271,7 +271,6 @@ func CreateTableForQRep(conn *pgx.Conn, suffix string, tableName string) error {
 		"geography_linestring geography(linestring)",
 		"geometry_polygon geometry(polygon)",
 		"geography_polygon geography(polygon)",
-		"nannu NUMERIC",
 		"myreal REAL",
 		"myreal2 REAL",
 		"myreal3 REAL",
@@ -337,12 +336,8 @@ func PopulateSourceTable(conn *pgx.Conn, suffix string, tableName string, rowCou
 						'LINESTRING(0 0, 1 1, 2 2)',
 						'LINESTRING(-74.0060 40.7128, -73.9352 40.7306, -73.9123 40.7831)',
 						'POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))','POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))',
-						'NaN',
-						3.14159,
-						1,
-						1.0,
-						'10.0.0.0/32',
-						'1.1.10.2'::cidr
+						3.14159, 1, 1.0,
+						'10.0.0.0/32', '1.1.10.2'::cidr
 					)`,
 			id, uuid.New().String(), uuid.New().String(),
 			uuid.New().String(), uuid.New().String(), uuid.New().String(), uuid.New().String())
@@ -360,12 +355,8 @@ func PopulateSourceTable(conn *pgx.Conn, suffix string, tableName string, rowCou
 					settle_at, settlement_delay_reason, f1, f2, f3, f4, f5, f6, f7, f8, f9, f10, f11, f12, f13, my_date,
 					my_time, my_mood, myh,
 					"geometryPoint", geography_point,geometry_linestring, geography_linestring,geometry_polygon, geography_polygon,
-					nannu,
-					myreal,
-					myreal2,
-					myreal3,
-					myinet,
-					mycidr
+					myreal, myreal2, myreal3,
+					myinet, mycidr
 			) VALUES %s;
 	`, suffix, tableName, strings.Join(rows, ",")))
 	if err != nil {