test cleanup

Author: mieciu

platform/frontend_connectors/schema_transformer_test.go

@@ -2154,87 +2154,98 @@ func Test_acceptIntsAsTimestamps(t *testing.T) {
 }
 
 func TestApplySelectFromCluster(t *testing.T) {
-	query := &model.Query{
-		TableName: "test_table",
-		SelectCommand: model.SelectCommand{
-			FromClause: model.NewTableRefWithDatabaseName("test_table", "test_db"),
-		},
-	}
+	originalTableReference := model.NewTableRefWithDatabaseName("test_table", "test_db")
+	expectedTableReference := model.NewLiteral(`cluster("my_cluster", "test_db", "test_table")`)
 	cfg := &config.QuesmaConfiguration{
 		ClusterName: "my_cluster",
 	}
-	transform := NewSchemaCheckPass(cfg, nil, defaultSearchAfterStrategy)
 
-	result, err := transform.ApplySelectFromCluster(schema.Schema{}, query)
-	assert.NoError(t, err)
-	expected := model.NewLiteral(`cluster("my_cluster", "test_db", "test_table")`)
-	assert.Equal(t, expected, result.SelectCommand.FromClause)
-}
-
-func TestApplySelectFromCluster2(t *testing.T) {
-
-	query := &model.Query{
-		SelectCommand: model.SelectCommand{
-			Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("FirstColumn"))},
-			FromClause: model.SelectCommand{
-				Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("SecondColumn"))},
-				FromClause: model.SelectCommand{
-					Columns:    []model.Expr{model.NewFunction("sum", model.NewColumnRef("ThirdColumn"))},
-					FromClause: model.NewTableRefWithDatabaseName("test_table", "test_db"),
+	tests := []struct {
+		name     string
+		query    *model.Query
+		expected *model.Query
+	}{
+		{
+			name: "simple FROM <tableName> substitution",
+			query: &model.Query{
+				TableName: "test_table",
+				SelectCommand: model.SelectCommand{
+					FromClause: originalTableReference,
+				},
+			},
+			expected: &model.Query{
+				TableName: "test_table",
+				SelectCommand: model.SelectCommand{
+					FromClause: expectedTableReference,
+				},
+			},
+		},
+		{
+			name: "nested FROM (SELECT ...) substitution",
+			query: &model.Query{
+				SelectCommand: model.SelectCommand{
+					Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("FirstColumn"))},
+					FromClause: model.SelectCommand{
+						Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("SecondColumn"))},
+						FromClause: model.SelectCommand{
+							Columns:    []model.Expr{model.NewFunction("sum", model.NewColumnRef("ThirdColumn"))},
+							FromClause: originalTableReference,
+							WhereClause: model.NewInfixExpr(
+								model.NewColumnRef("ThirdColumn"),
+								">=",
+								model.NewLiteral(50),
+							),
+						},
+						WhereClause: model.NewInfixExpr(
+							model.NewColumnRef("SecondColumn"),
+							">=",
+							model.NewLiteral(50),
+						),
+					},
 					WhereClause: model.NewInfixExpr(
-						model.NewColumnRef("ThirdColumn"),
+						model.NewColumnRef("FirstColumn"),
 						">=",
 						model.NewLiteral(50),
 					),
 				},
-				WhereClause: model.NewInfixExpr(
-					model.NewColumnRef("SecondColumn"),
-					">=",
-					model.NewLiteral(50),
-				),
-			},
-			WhereClause: model.NewInfixExpr(
-				model.NewColumnRef("FirstColumn"),
-				">=",
-				model.NewLiteral(50),
-			),
-		},
-	}
-	cfg := &config.QuesmaConfiguration{
-		ClusterName: "my_cluster",
-	}
-	transform := NewSchemaCheckPass(cfg, nil, defaultSearchAfterStrategy)
-
-	// When
-	result, err := transform.ApplySelectFromCluster(schema.Schema{}, query)
-	// Then
-	assert.NoError(t, err)
-	expected := &model.Query{
-		SelectCommand: model.SelectCommand{
-			Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("FirstColumn"))},
-			FromClause: model.SelectCommand{
-				Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("SecondColumn"))},
-				FromClause: model.SelectCommand{
-					Columns:    []model.Expr{model.NewFunction("sum", model.NewColumnRef("ThirdColumn"))},
-					FromClause: model.NewLiteral(`cluster("my_cluster", "test_db", "test_table")`),
+			},
+			expected: &model.Query{
+				SelectCommand: model.SelectCommand{
+					Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("FirstColumn"))},
+					FromClause: model.SelectCommand{
+						Columns: []model.Expr{model.NewFunction("sum", model.NewColumnRef("SecondColumn"))},
+						FromClause: model.SelectCommand{
+							Columns:    []model.Expr{model.NewFunction("sum", model.NewColumnRef("ThirdColumn"))},
+							FromClause: expectedTableReference,
+							WhereClause: model.NewInfixExpr(
+								model.NewColumnRef("ThirdColumn"),
+								">=",
+								model.NewLiteral(50),
+							),
+						},
+						WhereClause: model.NewInfixExpr(
+							model.NewColumnRef("SecondColumn"),
+							">=",
+							model.NewLiteral(50),
+						),
+					},
 					WhereClause: model.NewInfixExpr(
-						model.NewColumnRef("ThirdColumn"),
+						model.NewColumnRef("FirstColumn"),
 						">=",
 						model.NewLiteral(50),
 					),
 				},
-				WhereClause: model.NewInfixExpr(
-					model.NewColumnRef("SecondColumn"),
-					">=",
-					model.NewLiteral(50),
-				),
-			},
-			WhereClause: model.NewInfixExpr(
-				model.NewColumnRef("FirstColumn"),
-				">=",
-				model.NewLiteral(50),
-			),
+			},
 		},
 	}
-	assert.Equal(t, expected, result)
+
+	transform := NewSchemaCheckPass(cfg, nil, defaultSearchAfterStrategy)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := transform.ApplySelectFromCluster(schema.Schema{}, tt.query)
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
 }