add basic integration test, move to mcp

datahub-project · Jul 12, 2024 · 2e05fb0 · 2e05fb0
1 parent 3312b10
commit 2e05fb0
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 33 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/grafana/grafana_source.py b/metadata-ingestion/src/datahub/ingestion/source/grafana/grafana_source.py
@@ -25,7 +25,7 @@
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import ChangeAuditStamps
-from datahub.metadata.schema_classes import DashboardInfoClass
+from datahub.metadata.schema_classes import DashboardInfoClass, StatusClass
 
 
 class GrafanaSourceConfig(StatefulIngestionConfigBase, PlatformInstanceConfigMixin):
@@ -92,36 +92,37 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
             _title = item["title"]
             _url = item["url"]
             full_url = f"{self.source_config.url}{_url}"
-            _folder_id = item.get("folderId", None)
-            if _folder_id is not None:
-                dashboard_urn = builder.make_dashboard_urn(
-                    platform=self.platform,
-                    name=_uid,
-                    platform_instance=self.source_config.platform_instance,
+            dashboard_urn = builder.make_dashboard_urn(
+                platform=self.platform,
+                name=_uid,
+                platform_instance=self.source_config.platform_instance,
+            )
+            for mcp in MetadataChangeProposalWrapper.construct_many(
+                entityUrn=dashboard_urn,
+                aspects=[
+                    DashboardInfoClass(
+                        description="",
+                        title=_title,
+                        charts=[],
+                        lastModified=ChangeAuditStamps(),
+                        dashboardUrl=full_url,
+                        customProperties={
+                            "displayName": _title,
+                            "id": str(item["id"]),
+                            "uid": _uid,
+                            "title": _title,
+                            "uri": item["uri"],
+                            "type": item["type"],
+                            "folderId": str(item.get("folderId", None)),
+                            "folderUid": item.get("folderUid", None),
+                            "folderTitle": str(item.get("folderTitle", None)),
+                        },
+                    ),
+                    StatusClass(removed=False),
+                ],
+            ):
+                breakpoint()
+                yield MetadataWorkUnit(
+                    id=dashboard_urn,
+                    mcp=mcp,
                 )
-                yield from [
-                    mcp.as_workunit()
-                    for mcp in MetadataChangeProposalWrapper.construct_many(
-                        entityUrn=dashboard_urn,
-                        aspects=[
-                            DashboardInfoClass(
-                                description="",
-                                title=_title,
-                                charts=[],
-                                lastModified=ChangeAuditStamps(),
-                                dashboardUrl=full_url,
-                                customProperties={
-                                    "displayName": _title,
-                                    "id": str(item["id"]),
-                                    "uid": _uid,
-                                    "title": _title,
-                                    "uri": item["uri"],
-                                    "type": item["type"],
-                                    "folderId": str(item.get("folderId", None)),
-                                    "folderUid": item.get("folderUid", None),
-                                    "folderTitle": str(item.get("folderTitle", None)),
-                                },
-                            )
-                        ],
-                    )
-                ]
diff --git a/metadata-ingestion/tests/integration/grafana/default-dashboard.json b/metadata-ingestion/tests/integration/grafana/default-dashboard.json
@@ -0,0 +1,25 @@
+{
+    "id": null,
+    "uid": "default",
+    "title": "Default Dashboard",
+    "tags": [],
+    "timezone": "browser",
+    "schemaVersion": 16,
+    "version": 0,
+    "panels": [
+        {
+            "type": "text",
+            "title": "Welcome",
+            "gridPos": {
+                "x": 0,
+                "y": 0,
+                "w": 24,
+                "h": 5
+            },
+            "options": {
+                "content": "Welcome to your Grafana dashboard!",
+                "mode": "markdown"
+            }
+        }
+    ]
+}
diff --git a/metadata-ingestion/tests/integration/grafana/grafana_mcps_golden.json b/metadata-ingestion/tests/integration/grafana/grafana_mcps_golden.json
@@ -0,0 +1,18 @@
+[
+{
+    "entityType": "dashboard",
+    "entityUrn": "urn:li:dashboard:(grafana,default)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    },
+    "systemMetadata": {
+        "lastObserved": 1720785600000,
+        "runId": "grafana-test-simple",
+        "lastRunId": "no-run-id-provided"
+    }
+}
+]
diff --git a/metadata-ingestion/tests/integration/grafana/provisioning/api-keys/api_keys.yaml b/metadata-ingestion/tests/integration/grafana/provisioning/api-keys/api_keys.yaml
@@ -0,0 +1,3 @@
+api_keys:
+  - name: 'example-api-key'
+    role: 'Admin'
diff --git a/metadata-ingestion/tests/integration/grafana/provisioning/dashboards/dashboard.yaml b/metadata-ingestion/tests/integration/grafana/provisioning/dashboards/dashboard.yaml
@@ -0,0 +1,11 @@
+apiVersion: 1
+
+providers:
+  - name: 'default'
+    orgId: 1
+    folder: ''
+    type: file
+    disableDeletion: false
+    updateIntervalSeconds: 10
+    options:
+      path: /var/lib/grafana/dashboards
diff --git a/metadata-ingestion/tests/integration/grafana/provisioning/datasources/datasource.yaml b/metadata-ingestion/tests/integration/grafana/provisioning/datasources/datasource.yaml
@@ -0,0 +1,12 @@
+apiVersion: 1
+
+datasources:
+  - name: PostgreSQL
+    type: postgres
+    access: proxy
+    url: postgres:5432
+    database: grafana
+    user: grafana
+    password: grafana
+    jsonData:
+      sslmode: disable
diff --git a/...a-ingestion/tests/integration/grafana/provisioning/service_accounts/service_accounts.yaml b/...a-ingestion/tests/integration/grafana/provisioning/service_accounts/service_accounts.yaml
@@ -0,0 +1,6 @@
+service_accounts:
+  - name: 'example-service-account'
+    role: 'Admin'
+    apiKeys:
+      - keyName: 'example-api-key'
+        role: 'Admin'