Updated test with new crawler

FastLee · FastLee · commit a84381bab4c1 · 2023-10-04T16:35:11.000-04:00
diff --git a/src/databricks/labs/ucx/framework/crawlers.py b/src/databricks/labs/ucx/framework/crawlers.py
@@ -22,10 +22,10 @@ def fetch(self, sql) -> Iterator[any]:
         raise NotImplementedError
 
     @abstractmethod
-    def save_table(self, full_name: str, rows: list[any], klass: type, mode: str = "append"):
+    def save_table(self, full_name: str, rows: list[any], klass: dataclasses.dataclass, mode: str = "append"):
         raise NotImplementedError
 
-    def create_table(self, full_name: str, klass: type):
+    def create_table(self, full_name: str, klass: dataclasses.dataclass):
         ddl = f"CREATE TABLE IF NOT EXISTS {full_name} ({self._schema_for(klass)}) USING DELTA"
         self.execute(ddl)
 
@@ -86,7 +86,7 @@ def fetch(self, sql) -> Iterator[any]:
         logger.debug(f"[api][fetch] {sql}")
         return self._sql.execute_fetch_all(self._warehouse_id, sql)
 
-    def save_table(self, full_name: str, rows: list[any], klass: type, mode="append"):
+    def save_table(self, full_name: str, rows: list[any], klass: dataclasses.dataclass, mode="append"):
         if mode == "overwrite":
             msg = "Overwrite mode is not yet supported"
             raise NotImplementedError(msg)
@@ -140,7 +140,7 @@ def fetch(self, sql) -> Iterator[any]:
         logger.debug(f"[spark][fetch] {sql}")
         return self._spark.sql(sql).collect()
 
-    def save_table(self, full_name: str, rows: list[any], klass: type, mode: str = "append"):
+    def save_table(self, full_name: str, rows: list[any], klass: dataclasses.dataclass, mode: str = "append"):
         rows = self._filter_none_rows(rows, full_name)
 
         if len(rows) == 0:
@@ -152,7 +152,7 @@ def save_table(self, full_name: str, rows: list[any], klass: type, mode: str = "
 
 
 class CrawlerBase:
-    def __init__(self, backend: SqlBackend, catalog: str, schema: str, table: str, klass: type):
+    def __init__(self, backend: SqlBackend, catalog: str, schema: str, table: str, klass: dataclasses.dataclass):
         """
         Initializes a CrawlerBase instance.
 
diff --git a/tests/integration/hive_metastore/test_external_locations.py b/tests/integration/hive_metastore/test_external_locations.py
@@ -9,7 +9,7 @@
 logger = logging.getLogger(__name__)
 
 
-def test_table_inventory(ws, make_warehouse, make_schema):
+def test_external_locations(ws, make_warehouse, make_schema):
     warehouse_id = os.environ["TEST_DEFAULT_WAREHOUSE_ID"]
 
     logger.info("setting up fixtures")
@@ -20,8 +20,8 @@ def test_table_inventory(ws, make_warehouse, make_schema):
         Table("hive_metastore", "foo", "bar", "EXTERNAL", "delta", location="dbfs:/mnt/foo/test3/table3"),
     ]
     schema = make_schema()
-    sbe.save_table(f"{schema}.tables", tables)
-    sbe.save_table(f"{schema}.mounts", [Mount("/mnt/foo", "s3://bar")])
+    sbe.save_table(f"{schema}.tables", tables, Table)
+    sbe.save_table(f"{schema}.mounts", [Mount("/mnt/foo", "s3://bar")], Mount)
 
     crawler = ExternalLocationCrawler(ws, sbe, schema.split(".")[1])
     results = crawler.snapshot()
