Better mocking of os.environ

Signed-off-by: Thomas Newton <[email protected]>
flyteorg · Sep 25, 2023 · b689b3c · b689b3c
1 parent c3ee81f
commit b689b3c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 34 deletions.
diff --git a/tests/flytekit/unit/core/test_data.py b/tests/flytekit/unit/core/test_data.py
@@ -227,44 +227,41 @@ def test_s3_setup_args_env_aws(mock_os, mock_get_config_file):
 
 
 @mock.patch("flytekit.configuration.get_config_file")
-@mock.patch.dict(
-    os.environ,
-    {
-        "FLYTE_GCP_GSUTIL_PARALLELISM": "False",
-    },
-)
-def test_get_fsspec_storage_options_gcs(mock_get_config_file):
+@mock.patch("os.environ")
+def test_get_fsspec_storage_options_gcs(mock_os, mock_get_config_file):
     mock_get_config_file.return_value = None
+    ee = {
+        "FLYTE_GCP_GSUTIL_PARALLELISM": "False",
+    }
+    mock_os.get.side_effect = lambda x, y: ee.get(x)
     storage_options = get_fsspec_storage_options("gs", DataConfig.auto())
     assert storage_options == {}
 
 
 @mock.patch("flytekit.configuration.get_config_file")
-@mock.patch.dict(
-    os.environ,
-    {
-        "FLYTE_GCP_GSUTIL_PARALLELISM": "False",
-    },
-)
-def test_get_fsspec_storage_options_gcs_with_overrides(mock_get_config_file):
+@mock.patch("os.environ")
+def test_get_fsspec_storage_options_gcs_with_overrides(mock_os, mock_get_config_file):
     mock_get_config_file.return_value = None
+    ee = {
+        "FLYTE_GCP_GSUTIL_PARALLELISM": "False",
+    }
+    mock_os.get.side_effect = lambda x, y: ee.get(x)
     storage_options = get_fsspec_storage_options("gs", DataConfig.auto(), anonymous=True, other_argument="value")
     assert storage_options == {"token": "anon", "other_argument": "value"}
 
 
 @mock.patch("flytekit.configuration.get_config_file")
-@mock.patch.dict(
-    os.environ,
-    {
+@mock.patch("os.environ")
+def test_get_fsspec_storage_options_azure(mock_os, mock_get_config_file):
+    mock_get_config_file.return_value = None
+    ee = {
         "FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname",
         "FLYTE_AZURE_STORAGE_ACCOUNT_KEY": "accountkey",
         "FLYTE_AZURE_TENANT_ID": "tenantid",
         "FLYTE_AZURE_CLIENT_ID": "clientid",
         "FLYTE_AZURE_CLIENT_SECRET": "clientsecret",
-    },
-)
-def test_get_fsspec_storage_options_azure(mock_get_config_file):
-    mock_get_config_file.return_value = None
+    }
+    mock_os.get.side_effect = lambda x, y: ee.get(x)
     storage_options = get_fsspec_storage_options("abfs", DataConfig.auto())
     assert storage_options == {
         "account_name": "accountname",
@@ -277,15 +274,14 @@ def test_get_fsspec_storage_options_azure(mock_get_config_file):
 
 
 @mock.patch("flytekit.configuration.get_config_file")
-@mock.patch.dict(
-    os.environ,
-    {
+@mock.patch("os.environ")
+def test_get_fsspec_storage_options_azure_with_overrides(mock_os, mock_get_config_file):
+    mock_get_config_file.return_value = None
+    ee = {
         "FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname",
         "FLYTE_AZURE_STORAGE_ACCOUNT_KEY": "accountkey",
-    },
-)
-def test_get_fsspec_storage_options_azure_with_overrides(mock_get_config_file):
-    mock_get_config_file.return_value = None
+    }
+    mock_os.get.side_effect = lambda x, y: ee.get(x)
     storage_options = get_fsspec_storage_options(
         "abfs", DataConfig.auto(), anonymous=True, account_name="other_accountname", other_argument="value"
     )

diff --git a/tests/flytekit/unit/core/test_data_persistence.py b/tests/flytekit/unit/core/test_data_persistence.py
@@ -1,7 +1,7 @@
 import os
 
+import mock
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
-from mock import patch
 
 from flytekit.core.data_persistence import FileAccessProvider
 
@@ -22,10 +22,9 @@ def test_is_remote():
 
 
 def test_initialise_azure_file_provider_with_account_key():
-    with patch.dict(
+    with mock.patch.dict(
         os.environ,
         {"FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname", "FLYTE_AZURE_STORAGE_ACCOUNT_KEY": "accountkey"},
-        clear=True,
     ):
         fp = FileAccessProvider("/tmp", "abfs://container/path/within/container")
         assert fp.get_filesystem().account_name == "accountname"
@@ -34,15 +33,14 @@ def test_initialise_azure_file_provider_with_account_key():
 
 
 def test_initialise_azure_file_provider_with_service_principal():
-    with patch.dict(
+    with mock.patch.dict(
         os.environ,
         {
             "FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname",
             "FLYTE_AZURE_CLIENT_SECRET": "clientsecret",
             "FLYTE_AZURE_CLIENT_ID": "clientid",
             "FLYTE_AZURE_TENANT_ID": "tenantid",
         },
-        clear=True,
     ):
         fp = FileAccessProvider("/tmp", "abfs://container/path/within/container")
         assert fp.get_filesystem().account_name == "accountname"
@@ -53,7 +51,7 @@ def test_initialise_azure_file_provider_with_service_principal():
 
 
 def test_initialise_azure_file_provider_with_default_credential():
-    with patch.dict(os.environ, {"FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname"}, clear=True):
+    with mock.patch.dict(os.environ, {"FLYTE_AZURE_STORAGE_ACCOUNT_NAME": "accountname"}):
         fp = FileAccessProvider("/tmp", "abfs://container/path/within/container")
         assert fp.get_filesystem().account_name == "accountname"
         assert isinstance(fp.get_filesystem().sync_credential, DefaultAzureCredential)