Utilise USE instead of DESC in connection test (#935)

src/snowflake/cli/api/sql_execution.py

@@ -8,6 +8,7 @@
 from typing import Iterable, Optional, Tuple
 
 from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import (
     DatabaseNotProvidedError,
     SchemaNotProvidedError,
@@ -62,6 +63,9 @@ def _execute_query(self, query: str, **kwargs):
     def _execute_queries(self, queries: str, **kwargs):
         return list(self._execute_string(dedent(queries), **kwargs))
 
+    def use(self, object_type: ObjectType, name: str):
+        return self._execute_query(f"use {object_type.value.sf_name} {name}")
+
     @contextmanager
     def use_role(self, new_role: str):
         """

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

@@ -254,17 +254,13 @@ def test(
     om = ObjectManager()
     try:
         if conn.role:
-            om.use_role(new_role=conn.role)
+            om.use(object_type=ObjectType.ROLE, name=conn.role)
         if conn.database:
-            om.describe(
-                object_type=ObjectType.DATABASE.value.cli_name, name=conn.database
-            )
+            om.use(object_type=ObjectType.DATABASE, name=conn.database)
         if conn.schema:
-            om.describe(object_type=ObjectType.SCHEMA.value.cli_name, name=conn.schema)
+            om.use(object_type=ObjectType.SCHEMA, name=conn.schema)
         if conn.warehouse:
-            om.describe(
-                object_type=ObjectType.WAREHOUSE.value.cli_name, name=conn.warehouse
-            )
+            om.use(object_type=ObjectType.WAREHOUSE, name=conn.warehouse)
     except ProgrammingError as err:
         raise ClickException(str(err))
 

tests/test_connection.py

@@ -7,6 +7,7 @@
 
 import pytest
 import tomlkit
+from snowflake.cli.api.constants import ObjectType
 
 
 def test_new_connection_can_be_added(runner, snapshot):
@@ -303,11 +304,11 @@ def test_connection_test(mock_connect, mock_om, runner):
     )
 
     conn = mock_connect.return_value
-    mock_om.return_value.use_role.assert_called_once_with(new_role=conn.role)
-    assert mock_om.return_value.describe.mock_calls == [
-        mock.call(object_type="database", name=conn.database),
-        mock.call(object_type="schema", name=conn.schema),
-        mock.call(object_type="warehouse", name=conn.warehouse),
+    assert mock_om.return_value.use.mock_calls == [
+        mock.call(object_type=ObjectType.ROLE, name=conn.role),
+        mock.call(object_type=ObjectType.DATABASE, name=conn.database),
+        mock.call(object_type=ObjectType.SCHEMA, name=conn.schema),
+        mock.call(object_type=ObjectType.WAREHOUSE, name=conn.warehouse),
     ]
 
 

tests/test_sql.py

@@ -3,6 +3,7 @@
 from unittest import mock
 
 import pytest
+from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
 from snowflake.cli.api.project.util import identifier_to_show_like_pattern
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
@@ -253,3 +254,17 @@ def test_show_specific_object_multiple_rows(mock_execute_query):
     mock_execute_query.assert_called_once_with(
         r"show objects like 'NAME'", cursor_class=DictCursor
     )
+
+
+@pytest.mark.parametrize(
+    "_object",
+    [
+        ObjectType.WAREHOUSE,
+        ObjectType.ROLE,
+        ObjectType.DATABASE,
+    ],
+)
+@mock.patch("snowflake.cli.api.sql_execution.SqlExecutionMixin._execute_query")
+def test_use_command(mock_execute_query, _object):
+    SqlExecutionMixin().use(object_type=_object, name="foo_name")
+    mock_execute_query.assert_called_once_with(f"use {_object.value.sf_name} foo_name")