diff --git a/semantic_model_generator/tests/snowflake_connector_test.py b/semantic_model_generator/tests/snowflake_connector_test.py
index 774f7d44..f58e6bba 100644
--- 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,7 +232,7 @@ 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(
@@ -243,5 +240,5 @@ def test_fetch_valid_tables_and_views(snowflake_data, expected_df):
     )
 
     # 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")