NO-SNOW replace os.environ patching with monkeypatch everywhere in unit tests (#2500)

(cherry picked from commit d1d3880)

Author: sfc-gh-mmishchenko

test/unit/test_auth_workload_identity.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import os
 from base64 import b64decode
 from unittest import mock
 from urllib.parse import parse_qs, urlparse
@@ -414,11 +413,11 @@ def test_explicit_azure_omits_client_id_if_not_set(fake_azure_metadata_service):
     assert fake_azure_metadata_service.requested_client_id is None
 
 
-def test_explicit_azure_uses_explicit_client_id_if_set(fake_azure_metadata_service):
-    with mock.patch.dict(
-        os.environ, {"MANAGED_IDENTITY_CLIENT_ID": "custom-client-id"}
-    ):
-        auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AZURE)
-        auth_class.prepare(conn=None)
+def test_explicit_azure_uses_explicit_client_id_if_set(
+    fake_azure_metadata_service, monkeypatch
+):
+    monkeypatch.setenv("MANAGED_IDENTITY_CLIENT_ID", "custom-client-id")
+    auth_class = AuthByWorkloadIdentity(provider=AttestationProvider.AZURE)
+    auth_class.prepare(conn=None)
 
     assert fake_azure_metadata_service.requested_client_id == "custom-client-id"

test/unit/test_connection.py

@@ -3,7 +3,6 @@
 
 import json
 import logging
-import os
 import stat
 import sys
 from pathlib import Path
@@ -201,11 +200,11 @@ def test_is_still_running():
 
 
 @pytest.mark.skipolddriver
-def test_partner_env_var(mock_post_requests):
+def test_partner_env_var(mock_post_requests, monkeypatch):
     PARTNER_NAME = "Amanda"
 
-    with patch.dict(os.environ, {ENV_VAR_PARTNER: PARTNER_NAME}):
-        assert fake_connector().application == PARTNER_NAME
+    monkeypatch.setenv(ENV_VAR_PARTNER, PARTNER_NAME)
+    assert fake_connector().application == PARTNER_NAME
 
     assert (
         mock_post_requests["data"]["CLIENT_ENVIRONMENT"]["APPLICATION"] == PARTNER_NAME

test/unit/test_ocsp.py

@@ -10,7 +10,7 @@
 import platform
 import time
 from concurrent.futures.thread import ThreadPoolExecutor
-from os import environ, path
+from os import path
 from unittest import mock
 
 import asn1crypto.x509
@@ -78,7 +78,7 @@ def overwrite_ocsp_cache(tmpdir):
 
 
 @pytest.fixture(autouse=True)
-def worker_specific_cache_dir(tmpdir, request):
+def worker_specific_cache_dir(tmpdir, request, monkeypatch):
     """Create worker-specific cache directory to avoid file lock conflicts in parallel execution.
 
     Note: Tests that explicitly manage their own cache directories (like test_ocsp_cache_when_server_is_down)
@@ -88,13 +88,12 @@ def worker_specific_cache_dir(tmpdir, request):
     # Get worker ID for parallel execution (pytest-xdist)
     worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
 
-    # Store original cache dir environment variable
-    original_cache_dir = os.environ.get("SF_OCSP_RESPONSE_CACHE_DIR")
+    # monkeypatch will automatically handle restoration
 
     # Set worker-specific cache directory to prevent main cache file conflicts
     worker_cache_dir = tmpdir.join(f"ocsp_cache_{worker_id}")
     worker_cache_dir.ensure(dir=True)
-    os.environ["SF_OCSP_RESPONSE_CACHE_DIR"] = str(worker_cache_dir)
+    monkeypatch.setenv("SF_OCSP_RESPONSE_CACHE_DIR", str(worker_cache_dir))
 
     # Only handle the OCSP_RESPONSE_VALIDATION_CACHE to prevent conflicts
     # Let tests manage SF_OCSP_RESPONSE_CACHE_DIR themselves if they need to
@@ -131,11 +130,7 @@ def worker_specific_cache_dir(tmpdir, request):
         # If modules not available, just yield the directory
         yield str(tmpdir)
     finally:
-        # Restore original cache directory environment variable
-        if original_cache_dir is not None:
-            os.environ["SF_OCSP_RESPONSE_CACHE_DIR"] = original_cache_dir
-        else:
-            os.environ.pop("SF_OCSP_RESPONSE_CACHE_DIR", None)
+        # monkeypatch will automatically restore the original environment variable
 
         # Reset cache dir back to original state
         try:
@@ -235,7 +230,7 @@ def test_ocsp_wo_cache_server():
         assert ocsp.validate(url, connection), f"Failed to validate: {url}"
 
 
-def test_ocsp_wo_cache_file():
+def test_ocsp_wo_cache_file(monkeypatch):
     """OCSP tests without File cache.
 
     Notes:
@@ -248,7 +243,7 @@ def test_ocsp_wo_cache_file():
     except FileNotFoundError:
         # File doesn't exist, which is fine for this test
         pass
-    environ["SF_OCSP_RESPONSE_CACHE_DIR"] = "/etc"
+    monkeypatch.setenv("SF_OCSP_RESPONSE_CACHE_DIR", "/etc")
     OCSPCache.reset_cache_dir()
 
     try:
@@ -257,11 +252,10 @@ def test_ocsp_wo_cache_file():
             connection = _openssl_connect(url)
             assert ocsp.validate(url, connection), f"Failed to validate: {url}"
     finally:
-        del environ["SF_OCSP_RESPONSE_CACHE_DIR"]
         OCSPCache.reset_cache_dir()
 
 
-def test_ocsp_fail_open_w_single_endpoint():
+def test_ocsp_fail_open_w_single_endpoint(monkeypatch):
     SnowflakeOCSP.clear_cache()
 
     try:
@@ -270,33 +264,28 @@ def test_ocsp_fail_open_w_single_endpoint():
         # File doesn't exist, which is fine for this test
         pass
 
-    environ["SF_OCSP_TEST_MODE"] = "true"
-    environ["SF_TEST_OCSP_URL"] = "http://httpbin.org/delay/10"
-    environ["SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT"] = "5"
+    monkeypatch.setenv("SF_OCSP_TEST_MODE", "true")
+    monkeypatch.setenv("SF_TEST_OCSP_URL", "http://httpbin.org/delay/10")
+    monkeypatch.setenv("SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT", "5")
 
     ocsp = SFOCSP(use_ocsp_cache_server=False)
     connection = _openssl_connect("snowflake.okta.com")
 
-    try:
-        assert ocsp.validate(
-            "snowflake.okta.com", connection
-        ), "Failed to validate: {}".format("snowflake.okta.com")
-    finally:
-        del environ["SF_OCSP_TEST_MODE"]
-        del environ["SF_TEST_OCSP_URL"]
-        del environ["SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT"]
+    assert ocsp.validate(
+        "snowflake.okta.com", connection
+    ), "Failed to validate: {}".format("snowflake.okta.com")
 
 
 @pytest.mark.skipif(
     ER_OCSP_RESPONSE_CERT_STATUS_REVOKED is None,
     reason="No ER_OCSP_RESPONSE_CERT_STATUS_REVOKED is available.",
 )
-def test_ocsp_fail_close_w_single_endpoint():
+def test_ocsp_fail_close_w_single_endpoint(monkeypatch):
     SnowflakeOCSP.clear_cache()
 
-    environ["SF_OCSP_TEST_MODE"] = "true"
-    environ["SF_TEST_OCSP_URL"] = "http://httpbin.org/delay/10"
-    environ["SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT"] = "5"
+    monkeypatch.setenv("SF_OCSP_TEST_MODE", "true")
+    monkeypatch.setenv("SF_TEST_OCSP_URL", "http://httpbin.org/delay/10")
+    monkeypatch.setenv("SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT", "5")
 
     OCSPCache.del_cache_file()
 
@@ -306,21 +295,16 @@ def test_ocsp_fail_close_w_single_endpoint():
     with pytest.raises(RevocationCheckError) as ex:
         ocsp.validate("snowflake.okta.com", connection)
 
-    try:
-        assert (
-            ex.value.errno == ER_OCSP_RESPONSE_FETCH_FAILURE
-        ), "Connection should have failed"
-    finally:
-        del environ["SF_OCSP_TEST_MODE"]
-        del environ["SF_TEST_OCSP_URL"]
-        del environ["SF_TEST_CA_OCSP_RESPONDER_CONNECTION_TIMEOUT"]
+    assert (
+        ex.value.errno == ER_OCSP_RESPONSE_FETCH_FAILURE
+    ), "Connection should have failed"
 
 
-def test_ocsp_bad_validity():
+def test_ocsp_bad_validity(monkeypatch):
     SnowflakeOCSP.clear_cache()
 
-    environ["SF_OCSP_TEST_MODE"] = "true"
-    environ["SF_TEST_OCSP_FORCE_BAD_RESPONSE_VALIDITY"] = "true"
+    monkeypatch.setenv("SF_OCSP_TEST_MODE", "true")
+    monkeypatch.setenv("SF_TEST_OCSP_FORCE_BAD_RESPONSE_VALIDITY", "true")
 
     try:
         OCSPCache.del_cache_file()
@@ -334,12 +318,10 @@ def test_ocsp_bad_validity():
     assert ocsp.validate(
         "snowflake.okta.com", connection
     ), "Connection should have passed with fail open"
-    del environ["SF_OCSP_TEST_MODE"]
-    del environ["SF_TEST_OCSP_FORCE_BAD_RESPONSE_VALIDITY"]
 
 
-def test_ocsp_single_endpoint():
-    environ["SF_OCSP_ACTIVATE_NEW_ENDPOINT"] = "True"
+def test_ocsp_single_endpoint(monkeypatch):
+    monkeypatch.setenv("SF_OCSP_ACTIVATE_NEW_ENDPOINT", "True")
     SnowflakeOCSP.clear_cache()
     ocsp = SFOCSP()
     ocsp.OCSP_CACHE_SERVER.NEW_DEFAULT_CACHE_SERVER_BASE_URL = "https://snowflake.preprod3.us-west-2-dev.external-zone.snowflakecomputing.com:8085/ocsp/"
@@ -348,8 +330,6 @@ def test_ocsp_single_endpoint():
         "snowflake.okta.com", connection
     ), "Failed to validate: {}".format("snowflake.okta.com")
 
-    del environ["SF_OCSP_ACTIVATE_NEW_ENDPOINT"]
-
 
 def test_ocsp_by_post_method():
     """OCSP tests."""
@@ -375,15 +355,17 @@ def test_ocsp_with_file_cache(tmpdir):
 
 
 @pytest.mark.skipolddriver
-def test_ocsp_with_bogus_cache_files(tmpdir, random_ocsp_response_validation_cache):
+def test_ocsp_with_bogus_cache_files(
+    tmpdir, random_ocsp_response_validation_cache, monkeypatch
+):
     with mock.patch(
         "snowflake.connector.ocsp_snowflake.OCSP_RESPONSE_VALIDATION_CACHE",
         random_ocsp_response_validation_cache,
     ):
         from snowflake.connector.ocsp_snowflake import OCSPResponseValidationResult
 
         """Attempts to use bogus OCSP response data."""
-        cache_file_name, target_hosts = _store_cache_in_file(tmpdir)
+        cache_file_name, target_hosts = _store_cache_in_file(monkeypatch, tmpdir)
 
         ocsp = SFOCSP()
         OCSPCache.read_ocsp_response_cache_file(ocsp, cache_file_name)
@@ -414,15 +396,17 @@ def test_ocsp_with_bogus_cache_files(tmpdir, random_ocsp_response_validation_cac
 
 
 @pytest.mark.skipolddriver
-def test_ocsp_with_outdated_cache(tmpdir, random_ocsp_response_validation_cache):
+def test_ocsp_with_outdated_cache(
+    tmpdir, random_ocsp_response_validation_cache, monkeypatch
+):
     with mock.patch(
         "snowflake.connector.ocsp_snowflake.OCSP_RESPONSE_VALIDATION_CACHE",
         random_ocsp_response_validation_cache,
     ):
         from snowflake.connector.ocsp_snowflake import OCSPResponseValidationResult
 
         """Attempts to use outdated OCSP response cache file."""
-        cache_file_name, target_hosts = _store_cache_in_file(tmpdir)
+        cache_file_name, target_hosts = _store_cache_in_file(monkeypatch, tmpdir)
 
         ocsp = SFOCSP()
 
@@ -452,10 +436,8 @@ def test_ocsp_with_outdated_cache(tmpdir, random_ocsp_response_validation_cache)
         ), "must be empty. outdated cache should not be loaded"
 
 
-def _store_cache_in_file(tmpdir, target_hosts=None):
-    if target_hosts is None:
-        target_hosts = TARGET_HOSTS
-    os.environ["SF_OCSP_RESPONSE_CACHE_DIR"] = str(tmpdir)
+def _store_cache_in_file(monkeypatch, tmpdir):
+    monkeypatch.setenv("SF_OCSP_RESPONSE_CACHE_DIR", str(tmpdir))
     OCSPCache.reset_cache_dir()
     filename = path.join(str(tmpdir), "ocsp_response_cache.json")
 
@@ -464,13 +446,13 @@ def _store_cache_in_file(tmpdir, target_hosts=None):
     ocsp = SFOCSP(
         ocsp_response_cache_uri="file://" + filename, use_ocsp_cache_server=False
     )
-    for hostname in target_hosts:
+    for hostname in TARGET_HOSTS:
         connection = _openssl_connect(hostname)
         assert ocsp.validate(hostname, connection), "Failed to validate: {}".format(
             hostname
         )
     assert path.exists(filename), "OCSP response cache file"
-    return filename, target_hosts
+    return filename, TARGET_HOSTS
 
 
 def test_ocsp_with_invalid_cache_file():
@@ -658,11 +640,11 @@ def test_building_retry_url():
     assert OCSP_SERVER.OCSP_RETRY_URL is None
 
 
-def test_building_new_retry():
+def test_building_new_retry(monkeypatch):
     OCSP_SERVER = OCSPServer()
     OCSP_SERVER.OCSP_RETRY_URL = None
     hname = "a1.us-east-1.snowflakecomputing.com"
-    os.environ["SF_OCSP_ACTIVATE_NEW_ENDPOINT"] = "true"
+    monkeypatch.setenv("SF_OCSP_ACTIVATE_NEW_ENDPOINT", "true")
     OCSP_SERVER.reset_ocsp_endpoint(hname)
     assert (
         OCSP_SERVER.CACHE_SERVER_URL
@@ -698,8 +680,6 @@ def test_building_new_retry():
         == "https://ocspssd.snowflakecomputing.com/ocsp/retry"
     )
 
-    del os.environ["SF_OCSP_ACTIVATE_NEW_ENDPOINT"]
-
 
 @pytest.mark.parametrize(
     "hash_algorithm",

test/unit/test_proxies.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import logging
-import os
 import unittest.mock
 
 import pytest
@@ -28,10 +27,10 @@ def test_get_proxy_url():
 
 
 @pytest.mark.skipolddriver
-def test_socks_5_proxy_missing_proxy_header_attribute(caplog):
+def test_socks_5_proxy_missing_proxy_header_attribute(caplog, monkeypatch):
     from snowflake.connector.vendored.urllib3.poolmanager import ProxyManager
 
-    os.environ["HTTPS_PROXY"] = "socks5://localhost:8080"
+    monkeypatch.setenv("HTTPS_PROXY", "socks5://localhost:8080")
 
     class MockSOCKSProxyManager:
         def __init__(self):
@@ -81,8 +80,6 @@ def mock_proxy_manager_for_url_wiht_header(*args, **kwargs):
             )
     assert "Unable to set 'Host' to proxy manager of type" not in caplog.text
 
-    del os.environ["HTTPS_PROXY"]
-
 
 @pytest.mark.skipolddriver
 @pytest.mark.parametrize("proxy_method", ["explicit_args", "env_vars"])