fix tests

Snowflake-Labs · Oct 14, 2024 · dcbd336 · dcbd336
1 parent 41afc01
commit dcbd336
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/semantic_model_generator/tests/snowflake_connector_test.py b/semantic_model_generator/tests/snowflake_connector_test.py
@@ -114,7 +114,6 @@ def test_connect(
 
     conn.cursor().execute.assert_has_calls(
         [
-            call("USE DATABASE test"),
             call("ALTER SESSION SET QUERY_TAG = 'SEMANTIC_MODEL_GENERATOR'"),
             call("ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = 120"),
         ]
@@ -138,8 +137,6 @@ def test_connect_with_schema(
 
     conn.cursor().execute.assert_has_calls(
         [
-            call("USE DATABASE test_db"),
-            call("USE SCHEMA test_schema"),
             call("ALTER SESSION SET QUERY_TAG = 'SEMANTIC_MODEL_GENERATOR'"),
             call("ALTER SESSION SET STATEMENT_TIMEOUT_IN_SECONDS = 120"),
         ]
@@ -235,13 +232,13 @@ def test_fetch_valid_tables_and_views(snowflake_data, expected_df):
     mock_cursor.description = mocked_descriptions
 
     # Call the function to test
-    result_df = snowflake_connector._fetch_valid_tables_and_views(mock_conn)
+    result_df = snowflake_connector._fetch_valid_tables_and_views(mock_conn, "mock_db")
 
     # Assert the result is as expected
     pd.testing.assert_frame_equal(
         result_df.reset_index(drop=True), expected_df.reset_index(drop=True)
     )
 
     # Verify execute was called with correct queries
-    mock_cursor.execute.assert_any_call("show tables in database")
-    mock_cursor.execute.assert_any_call("show views in database")
+    mock_cursor.execute.assert_any_call("show tables in database mock_db")
+    mock_cursor.execute.assert_any_call("show views in database mock_db")