From 4cd069759de61dd9367f0757a00ddbc5ecc862af Mon Sep 17 00:00:00 2001
From: spinillos <selenepinillos@gmail.com>
Date: Thu, 19 Sep 2024 18:04:43 +0200
Subject: [PATCH] Add compiler pass test

---
 .../set_datasource_to_dataquery_test.go       | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 internal/ast/compiler/set_datasource_to_dataquery_test.go

diff --git a/internal/ast/compiler/set_datasource_to_dataquery_test.go b/internal/ast/compiler/set_datasource_to_dataquery_test.go
new file mode 100644
index 00000000..29836d08
--- /dev/null
+++ b/internal/ast/compiler/set_datasource_to_dataquery_test.go
@@ -0,0 +1,66 @@
+package compiler
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/grafana/cog/internal/ast"
+	"github.com/grafana/cog/internal/testutils"
+)
+
+func TestSetDatasourceToDataQuery(t *testing.T) {
+	schemas := ast.Schemas{
+		&ast.Schema{
+			Package: "dashboard",
+			Objects: testutils.ObjectsMap(
+				ast.NewObject(dashboardPackage, dashboardDatasource, ast.NewStruct(
+					ast.NewStructField("Uid", ast.NewScalar(ast.KindString, ast.Nullable())),
+					ast.NewStructField("Type", ast.NewScalar(ast.KindString, ast.Nullable())),
+				)),
+			),
+		},
+		&ast.Schema{
+			Package: "no_dataquery",
+			Objects: testutils.ObjectsMap(
+				ast.NewObject("no_dataquery", datasourceName, ast.String()),
+				ast.NewObject("no_dataquery", dashboardDatasource, ast.String()),
+			),
+		},
+		&ast.Schema{
+			Package: "dataquery",
+			Metadata: ast.SchemaMeta{
+				Variant: ast.SchemaVariantDataQuery,
+			},
+			Objects: testutils.ObjectsMap(
+				ast.NewObject("dataquery", datasourceName, ast.Any()),
+				ast.NewObject("dataquery", "Query", ast.NewStruct(
+					ast.NewStructField("datasource", ast.NewRef("dataquery", datasourceName)),
+					ast.NewStructField("Type", ast.NewScalar(ast.KindString, ast.Nullable())),
+				)),
+			),
+		},
+	}
+
+	expected := ast.Schemas{
+		schemas[0],
+		schemas[1],
+		&ast.Schema{
+			Package: "dataquery",
+			Metadata: ast.SchemaMeta{
+				Variant: ast.SchemaVariantDataQuery,
+			},
+			Objects: testutils.ObjectsMap(
+				ast.NewObject("dataquery", "Query", ast.NewStruct(
+					ast.NewStructField(
+						"datasource",
+						ast.NewRef(dashboardPackage, dashboardDatasource),
+						ast.PassesTrail(fmt.Sprintf("SetDatasourceToDataquery[%s.%s]", dashboardPackage, dashboardDatasource)),
+					),
+					ast.NewStructField("Type", ast.NewScalar(ast.KindString, ast.Nullable())),
+				)),
+			),
+		},
+	}
+
+	runPassOnSchemas(t, &SetDatasourceToDataquery{}, schemas, expected)
+}