Snow 1313805 fix connection test

RELEASE-NOTES.md

@@ -53,6 +53,7 @@
 * Adding support for any source supported by `pip` in `snow snowpark`.
 * Fixed version parsing for packages lookup on Snowflake Anaconda Channel
 * Fix handling database/schema/role identifiers containing dashes
+* Fix schema override bug in `snow connection test`
 
 # v2.1.2
 

src/snowflake/cli/plugins/connection/commands.py

@@ -256,14 +256,19 @@ def test(
     # Test session attributes
     om = ObjectManager()
     try:
+        # "use database" operation changes schema to default "public",
+        # so to test schema set up by user we need to copy it here:
+        schema = conn.schema
+
         if conn.role:
             om.use(object_type=ObjectType.ROLE, name=f'"{conn.role}"')
         if conn.database:
             om.use(object_type=ObjectType.DATABASE, name=f'"{conn.database}"')
-        if conn.schema:
-            om.use(object_type=ObjectType.SCHEMA, name=f'"{conn.schema}"')
+        if schema:
+            om.use(object_type=ObjectType.SCHEMA, name=f'"{schema}"')
         if conn.warehouse:
             om.use(object_type=ObjectType.WAREHOUSE, name=f'"{conn.warehouse}"')
+
     except ProgrammingError as err:
         raise ClickException(str(err))
 

tests_integration/snowflake_connector.py

@@ -38,7 +38,7 @@ def setup_test_database(snowflake_session, database_name: str):
     snowflake_session.execute_string(
         f"create database {database_name}; use database {database_name}; use schema public;"
     )
-    with mock_single_env_var(DATABASE_ENV_PARAMETER, database_name):
+    with mock_single_env_var(DATABASE_ENV_PARAMETER, value=database_name):
         yield
     snowflake_session.execute_string(f"drop database {database_name}")
 
@@ -49,7 +49,7 @@ def setup_test_schema(snowflake_session, schema_name: str):
     snowflake_session.execute_string(
         f"create schema {schema_name}; use schema {schema_name};"
     )
-    with mock_single_env_var(SCHEMA_ENV_PARAMETER, schema_name):
+    with mock_single_env_var(SCHEMA_ENV_PARAMETER, value=schema_name):
         yield
     snowflake_session.execute_string(f"drop schema {schema_name}")
 

tests_integration/test_connection.py

@@ -3,6 +3,8 @@
     setup_test_database,
     setup_test_schema,
     add_uuid_to_name,
+    mock_single_env_var,
+    SCHEMA_ENV_PARAMETER,
 )
 
 
@@ -23,9 +25,6 @@ def test_connection_dashed_database(runner, snowflake_session):
 
 
 @pytest.mark.integration
-@pytest.mark.skip(
-    reason="BUG: connections test command seem to override every setting with schema PUBLIC"
-)
 def test_connection_dashed_schema(
     runner, test_database, snowflake_session, snowflake_home
 ):
@@ -34,3 +33,14 @@ def test_connection_dashed_schema(
         result = runner.invoke_with_connection(["connection", "test", "--debug"])
         assert result.exit_code == 0, result.output
         assert f'use schema "{schema}"' in result.output
+
+
+@pytest.mark.integration
+def test_connection_not_existing_schema(
+    runner, test_database, snowflake_session, snowflake_home
+):
+    schema = "schema_which_does_not_exist"
+    with mock_single_env_var(SCHEMA_ENV_PARAMETER, value=schema):
+        result = runner.invoke_with_connection(["connection", "test"])
+        assert result.exit_code == 1, result.output
+        assert "Object does not exist" in result.output