changing RunIntQuery to avoid casts

PeerDB-io · Dec 22, 2023 · 64906a2 · 64906a2
1 parent bc4a2ca
commit 64906a2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/flow/e2e/bigquery/bigquery_helper.go b/flow/e2e/bigquery/bigquery_helper.go
@@ -453,7 +453,7 @@ func (b *BigQueryTestHelper) CreateTable(tableName string, schema *model.QRecord
 	return nil
 }
 
-func (b *BigQueryTestHelper) RunIntQuery(query string) (int, error) {
+func (b *BigQueryTestHelper) RunInt64Query(query string) (int64, error) {
 	recordBatch, err := b.ExecuteAndProcessQuery(query)
 	if err != nil {
 		return 0, fmt.Errorf("could not execute query: %w", err)
@@ -462,5 +462,5 @@ func (b *BigQueryTestHelper) RunIntQuery(query string) (int, error) {
 		return 0, fmt.Errorf("expected only 1 record, got %d", recordBatch.NumRecords)
 	}
 
-	return int(recordBatch.Records[0].Entries[0].Value.(int64)), nil
+	return recordBatch.Records[0].Entries[0].Value.(int64), nil
 }
diff --git a/flow/e2e/bigquery/peer_flow_bq_test.go b/flow/e2e/bigquery/peer_flow_bq_test.go
@@ -1289,9 +1289,9 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_Basic() {
 	newerSyncedAtQuery := fmt.Sprintf(`
 		SELECT COUNT(*) FROM`+"`%s.%s`"+`WHERE _PEERDB_IS_DELETED = TRUE`,
 		s.bqHelper.datasetName, dstTableName)
-	numNewRows, err := s.bqHelper.RunIntQuery(newerSyncedAtQuery)
+	numNewRows, err := s.bqHelper.RunInt64Query(newerSyncedAtQuery)
 	require.NoError(s.t, err)
-	require.Equal(s.t, 1, numNewRows)
+	s.Eq(1, numNewRows)
 }
 
 func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_IUD_Same_Batch() {
@@ -1373,9 +1373,9 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_IUD_Same_Batch() {
 	newerSyncedAtQuery := fmt.Sprintf(`
 		SELECT COUNT(*) FROM`+"`%s.%s`"+`WHERE _PEERDB_IS_DELETED = TRUE`,
 		s.bqHelper.datasetName, dstTableName)
-	numNewRows, err := s.bqHelper.RunIntQuery(newerSyncedAtQuery)
+	numNewRows, err := s.bqHelper.RunInt64Query(newerSyncedAtQuery)
 	require.NoError(s.t, err)
-	s.Equal(1, numNewRows)
+	s.Eq(1, numNewRows)
 }
 
 func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_UD_Same_Batch() {
@@ -1461,9 +1461,9 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_UD_Same_Batch() {
 	newerSyncedAtQuery := fmt.Sprintf(`
 		SELECT COUNT(*) FROM`+"`%s.%s`"+`WHERE _PEERDB_IS_DELETED = TRUE`,
 		s.bqHelper.datasetName, dstTableName)
-	numNewRows, err := s.bqHelper.RunIntQuery(newerSyncedAtQuery)
+	numNewRows, err := s.bqHelper.RunInt64Query(newerSyncedAtQuery)
 	require.NoError(s.t, err)
-	s.Equal(1, numNewRows)
+	s.Eq(1, numNewRows)
 }
 
 func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_Insert_After_Delete() {
@@ -1537,7 +1537,7 @@ func (s PeerFlowE2ETestSuiteBQ) Test_Soft_Delete_Insert_After_Delete() {
 	newerSyncedAtQuery := fmt.Sprintf(`
 		SELECT COUNT(*) FROM`+"`%s.%s`"+`WHERE _PEERDB_IS_DELETED = TRUE`,
 		s.bqHelper.datasetName, dstTableName)
-	numNewRows, err := s.bqHelper.RunIntQuery(newerSyncedAtQuery)
+	numNewRows, err := s.bqHelper.RunInt64Query(newerSyncedAtQuery)
 	require.NoError(s.t, err)
-	s.Equal(0, numNewRows)
+	s.Eq(0, numNewRows)
 }