Fix tests

astronomer · Jan 30, 2024 · dc2e541 · dc2e541
1 parent 19824b7
commit dc2e541
Showing 1 changed file with 9 additions and 18 deletions.
diff --git a/tests/snowflake/hooks/test_snowflake_sql_api.py b/tests/snowflake/hooks/test_snowflake_sql_api.py
@@ -311,7 +311,7 @@ def test_get_private_key_should_support_private_auth_in_connection(
             hook.get_private_key()
             assert hook.private_key is not None
 
-    def test_get_private_key_raise_exception(self, encrypted_temporary_private_key: Path):
+    def test_get_private_key_raise_exception(self, encrypted_temporary_private_key: Path, monkeypatch):
         """
         Test get_private_key function with private_key_content and private_key_file in connection
         and raise airflow exception
@@ -329,17 +329,12 @@ def test_get_private_key_raise_exception(self, encrypted_temporary_private_key:
                 "private_key_file": str(encrypted_temporary_private_key),
             },
         }
+        monkeypatch.setenv("AIRFLOW_CONN_TEST_CONN", Connection(**connection_kwargs).get_uri())
         hook = SnowflakeSqlApiHookAsync(snowflake_conn_id="test_conn")
-        with (
-            unittest.mock.patch.dict(
-                "os.environ",
-                AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri(),
-            ),
-            pytest.raises(
-                AirflowException,
-                match="The private_key_file and private_key_content extra fields are mutually "
-                "exclusive. Please remove one.",
-            ),
+        with pytest.raises(
+            AirflowException,
+            match="The private_key_file and private_key_content extra fields are mutually "
+            "exclusive. Please remove one.",
         ):
             hook.get_private_key()
 
@@ -370,6 +365,7 @@ def test_get_private_key_should_support_private_auth_with_encrypted_key(
     def test_get_private_key_should_support_private_auth_with_unencrypted_key(
         self,
         non_encrypted_temporary_private_key,
+        monkeypatch,
     ):
         connection_kwargs = {
             **BASE_CONNECTION_KWARGS,
@@ -399,13 +395,8 @@ def test_get_private_key_should_support_private_auth_with_unencrypted_key(
             hook.get_private_key()
             assert hook.private_key is not None
         connection_kwargs["password"] = _PASSWORD
-        with (
-            unittest.mock.patch.dict(
-                "os.environ",
-                AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri(),
-            ),
-            pytest.raises(TypeError, match="Password was given but private key is not encrypted."),
-        ):
+        monkeypatch.setenv("AIRFLOW_CONN_TEST_CONN", Connection(**connection_kwargs).get_uri())
+        with pytest.raises(TypeError, match="Password was given but private key is not encrypted."):
             SnowflakeSqlApiHookAsync(snowflake_conn_id="test_conn").get_private_key()
 
     @pytest.mark.asyncio