fix test_missing_count

Author: haritamar

integration_tests/tests/data_seeder.py

@@ -1,6 +1,7 @@
 import csv
+from contextlib import contextmanager
 from pathlib import Path
-from typing import List
+from typing import Generator, List
 
 from elementary.clients.dbt.base_dbt_runner import BaseDbtRunner
 from logger import get_logger
@@ -18,7 +19,8 @@ def __init__(
         self.dbt_project_path = dbt_project_path
         self.seeds_dir_path = seeds_dir_path
 
-    def seed(self, data: List[dict], table_name: str):
+    @contextmanager
+    def seed(self, data: List[dict], table_name: str) -> Generator[None, None, None]:
         seed_path = self.seeds_dir_path.joinpath(f"{table_name}.csv")
         try:
             with seed_path.open("w") as seed_file:
@@ -28,5 +30,7 @@ def seed(self, data: List[dict], table_name: str):
                 writer.writerows(data)
                 seed_file.flush()
                 self.dbt_runner.seed(select=str(relative_seed_path), full_refresh=True)
+
+                yield
         finally:
             seed_path.unlink()

integration_tests/tests/dbt_project.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager, nullcontext
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from typing import Any, Dict, List, Literal, Optional, Union, overload
+from typing import Any, Dict, Generator, List, Literal, Optional, Union, overload
 from uuid import uuid4
 
 from data_seeder import DbtDataSeeder
@@ -235,9 +235,19 @@ def test(
             return [test_result] if multiple_results else test_result
 
     def seed(self, data: List[dict], table_name: str):
-        return DbtDataSeeder(
+        with DbtDataSeeder(
             self.dbt_runner, self.project_dir_path, self.seeds_dir_path
-        ).seed(data, table_name)
+        ).seed(data, table_name):
+            return
+
+    @contextmanager
+    def seed_context(
+        self, data: List[dict], table_name: str
+    ) -> Generator[None, None, None]:
+        with DbtDataSeeder(
+            self.dbt_runner, self.project_dir_path, self.seeds_dir_path
+        ).seed(data, table_name):
+            yield
 
     @contextmanager
     def create_temp_model_for_existing_table(

integration_tests/tests/test_string_monitors.py

@@ -6,9 +6,9 @@
 def test_missing_count(dbt_project: DbtProject, test_id: str):
     missing_values = [None, " ", "null", "NULL"]
     data = [{COLUMN_NAME: value} for value in ["a", "b", "c", " a "] + missing_values]
-    dbt_project.seed(data, test_id)
-    result = dbt_project.run_query(
-        f"select {{{{ elementary.missing_count('{COLUMN_NAME}') }}}} "
-        f"as missing_count from {{{{ generate_schema_name() }}}}.{test_id}"
-    )[0]
+    with dbt_project.seed_context(data, test_id):
+        result = dbt_project.run_query(
+            f"select {{{{ elementary.missing_count('{COLUMN_NAME}') }}}} "
+            f'as missing_count from {{{{ ref("{test_id}") }}}}'
+        )[0]
     assert result["missing_count"] == len(missing_values)