more test

src/databricks/labs/ucx/source_code/redash.py

@@ -18,41 +18,38 @@ def __init__(self, fixer: Fixer, ws: WorkspaceClient):
         self._fixer = fixer
         self._ws = ws
 
-    def fix_all_dashboards(self):
-        for dashboard in self._ws.dashboards.list():
-            self.fix_dashboard(dashboard)
-
-    def fix_dashboard(self, dashboard: Dashboard):
-        if dashboard.tags is not None and self.MIGRATED_TAG in dashboard.tags:
-            # already migrated
-            return
-        for query in self._get_queries_from_dashboard(dashboard):
-            self._fix_query(query)
-
-    def revert_dashboard(self, dashboard: Dashboard):
-        if dashboard.tags is None or self.MIGRATED_TAG not in dashboard.tags:
-            logger.debug(f"Dashboard {dashboard.id} was not migrated by UCX")
-            return
-        for query in self._get_queries_from_dashboard(dashboard):
-            self._revert_query(query)
-
-    def delete_backup_dashboards(self):
-        for dashboard in self._ws.dashboards.list():
-            if dashboard.tags is None or self.BACKUP_TAG not in dashboard.tags:
-                continue
-
+    def fix_dashboard(self, dashboard_id: str | None = None):
+        if dashboard_id is None:
+            dashboards = list(self._ws.dashboards.list())
+        else:
+            dashboards = [Dashboard.from_dict(self._get_dashboard(dashboard_id))]
+        for dashboard in dashboards:
+            if dashboard.tags is not None and self.MIGRATED_TAG in dashboard.tags:
+                # already migrated
+                return
             for query in self._get_queries_from_dashboard(dashboard):
-                if query.tags is None or self.BACKUP_TAG not in query.tags:
-                    continue
-                try:
-                    self._ws.queries.delete(query.id)
-                except NotFound:
-                    continue
+                self._fix_query(query)
+
+    def revert_dashboard(self, dashboard_id: str | None = None):
+        if dashboard_id is None:
+            dashboards = list(self._ws.dashboards.list())
+        else:
+            dashboards = [Dashboard.from_dict(self._get_dashboard(dashboard_id))]
+        for dashboard in dashboards:
+            if dashboard.tags is None or self.MIGRATED_TAG not in dashboard.tags:
+                logger.debug(f"Dashboard {dashboard.id} was not migrated by UCX")
+                return
+            for query in self._get_queries_from_dashboard(dashboard):
+                self._revert_query(query)
 
+    def delete_backup_dbsql_queries(self):
+        for query in self._ws.queries.list():
+            if query.tags is None or self.BACKUP_TAG not in query.tags:
+                continue
             try:
-                self._ws.dashboards.delete(dashboard.id)
+                self._ws.queries.delete(query.id)
             except NotFound:
-                pass
+                continue
 
     def _fix_query(self, query: Query):
         assert query.id is not None
@@ -116,13 +113,17 @@ def _revert_query(self, query: Query):
         self._ws.queries.delete(backup_id)
 
     def _get_queries_from_dashboard(self, dashboard: Dashboard) -> Iterator[Query]:
-        raw = self._ws.api_client.do("GET", f"/api/2.0/preview/sql/dashboards/{dashboard.id}")
-        if not isinstance(raw, dict):
-            return
+        raw = self._get_dashboard(dashboard.id)
         for widget in raw["widgets"]:
             if "visualization" not in widget:
                 continue
             viz = widget["visualization"]
             if "query" not in viz:
                 continue
             yield Query.from_dict(viz["query"])
+
+    def _get_dashboard(self, dashboard_id) -> dict:
+        raw = self._ws.api_client.do("GET", f"/api/2.0/preview/sql/dashboards/{dashboard_id}")
+        if not isinstance(raw, dict):
+            return {}
+        return raw

tests/integration/source_code/test_redash.py

@@ -4,7 +4,7 @@
 def test_fix_dashboard(ws, runtime_ctx, make_dashboard, make_query):
     dashboard, query = make_dashboard()
     another_query = make_query()
-    runtime_ctx.redash.fix_dashboard(dashboard)
+    runtime_ctx.redash.fix_dashboard(dashboard.id)
     # make sure the query is marked as migrated
     query = ws.queries.get(query.id)
     assert Redash.MIGRATED_TAG in query.tags
@@ -14,6 +14,6 @@ def test_fix_dashboard(ws, runtime_ctx, make_dashboard, make_query):
     assert len(another_query.tags) == 0
 
     # revert the dashboard, make sure the query has no tag
-    runtime_ctx.redash.revert_dashboard(dashboard)
+    runtime_ctx.redash.revert_dashboard(dashboard.id)
     query = ws.queries.get(query.id)
     assert len(query.tags) == 0

tests/unit/source_code/test_redash.py

@@ -1,4 +1,3 @@
-from unittest import mock
 from unittest.mock import create_autospec, call
 
 import pytest
@@ -22,6 +21,8 @@ def api_client_do(method: str, path: str, body: dict | None = None):
             dashboard_id = path.removeprefix("/api/2.0/preview/sql/dashboards/")
             tags = tags_lookup[dashboard_id]
             return {
+                "id": dashboard_id,
+                "tags": tags,
                 "widgets": [
                     {},
                     {"visualization": {}},
@@ -45,7 +46,7 @@ def api_client_do(method: str, path: str, body: dict | None = None):
                             }
                         }
                     },
-                ]
+                ],
             }
         if path.startswith("/api/2.0/preview/sql/queries"):
             return {}
@@ -70,7 +71,7 @@ def empty_fixer(empty_index):
 
 def test_fix_all_dashboards(redash_ws, empty_fixer):
     redash = Redash(empty_fixer, redash_ws)
-    redash.fix_all_dashboards()
+    redash.fix_dashboard()
     redash_ws.queries.create.assert_called_with(
         name='test_query_original',
         query='SELECT * FROM old.things',
@@ -82,7 +83,7 @@ def test_fix_all_dashboards(redash_ws, empty_fixer):
     )
     redash_ws.api_client.do.assert_has_calls(
         [
-            call('GET', '/api/2.0/preview/sql/dashboards/3'),
+            call('GET', '/api/2.0/preview/sql/dashboards/1'),
             # tag the backup query
             call('POST', '/api/2.0/preview/sql/queries/123', body={'tags': [Redash.BACKUP_TAG]}),
             # tag the migrated query
@@ -96,23 +97,26 @@ def test_fix_all_dashboards(redash_ws, empty_fixer):
 
 
 def test_revert_dashboard(redash_ws, empty_fixer):
+    redash_ws.queries.get.return_value = Query(id="1", query="original_query")
     redash = Redash(empty_fixer, redash_ws)
-    redash.revert_dashboard(Dashboard(id="2", tags=[Redash.MIGRATED_TAG]))
+    redash.revert_dashboard("2")
     redash_ws.api_client.do.assert_has_calls(
         [
+            call('GET', '/api/2.0/preview/sql/dashboards/2'),
             call('GET', '/api/2.0/preview/sql/dashboards/2'),
             # update the migrated query
             call(
                 'POST',
                 '/api/2.0/preview/sql/queries/1',
-                body={'query': mock.ANY, 'tags': []},
+                body={'query': 'original_query', 'tags': ['original']},
             ),
         ]
     )
     redash_ws.queries.delete.assert_called_once_with("123")
 
 
 def test_delete_backup_dashboards(redash_ws, empty_fixer):
+    redash_ws.queries.list.return_value = [Query(id="1", tags=[Redash.BACKUP_TAG]), Query(id="2", tags=[])]
     redash = Redash(empty_fixer, redash_ws)
-    redash.delete_backup_dashboards()
-    redash_ws.dashboards.delete.assert_called_once_with("1")
+    redash.delete_backup_dbsql_queries()
+    redash_ws.queries.delete.assert_called_once_with("1")