fix tests

pinterest · Oct 24, 2023 · 9ab9536 · 9ab9536
1 parent 57c7698
commit 9ab9536
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/querybook/server/lib/query_analysis/validation/decorators/metadata_decorators.py b/querybook/server/lib/query_analysis/validation/decorators/metadata_decorators.py
@@ -10,7 +10,7 @@
 from lib.query_analysis.validation.decorators.base_sqlglot_validation_decorator import (
     BaseValidationDecorator,
 )
-from logic.admin import get_query_engine_by_id, get_query_metastore_id_by_engine_id
+from logic import admin as admin_logic
 
 
 class BaseColumnNameSuggester(BaseValidationDecorator):
@@ -23,7 +23,7 @@ def get_column_name_from_error(
         raise NotImplementedError()
 
     def _get_tables_in_query(self, query: str, engine_id: int) -> List[str]:
-        engine = get_query_engine_by_id(engine_id)
+        engine = admin_logic.get_query_engine_by_id(engine_id)
         tables_per_statement, _ = process_query(query, language=engine.language)
         return list(chain.from_iterable(tables_per_statement))
 
@@ -91,7 +91,7 @@ def _suggest_table_name_if_needed(
         fuzzy_table_name = self.get_full_table_name_from_error(validation_result)
         if not fuzzy_table_name:
             return
-        metastore_id = get_query_metastore_id_by_engine_id(engine_id)
+        metastore_id = admin_logic.get_query_metastore_id_by_engine_id(engine_id)
         if metastore_id is None:
             return
         results, count = search_table.get_table_name_suggestion(

diff --git a/...b/test_query_analysis/test_validation/test_validators/test_presto_optimizing_validator.py b/...b/test_query_analysis/test_validation/test_validators/test_presto_optimizing_validator.py
@@ -373,6 +373,12 @@ def test__get_column_name_suggestion(self, mock_get_column_name_suggestion):
 class PrestoTableNameSuggesterTestCase(BaseValidatorTestCase):
     def setUp(self):
         self._validator = PrestoTableNameSuggester(MagicMock())
+        patch_get_metastore_id = patch(
+            "logic.admin.get_query_metastore_id_by_engine_id"
+        )
+        mock_get_metastore_id = patch_get_metastore_id.start()
+        mock_get_metastore_id.return_value = 1
+        self.addCleanup(patch_get_metastore_id.stop)
 
     def test_get_full_table_name_from_error(self):
         self.assertEquals(