From b689b3cd5159244b60ec3a82e454bbb775bf85c6 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Mon, 25 Sep 2023 14:55:17 +0100 Subject: [PATCH] Better mocking of os.environ Signed-off-by: Thomas Newton --- tests/flytekit/unit/core/test_data.py | 52 +++++++++---------- .../unit/core/test_data_persistence.py | 10 ++-- 2 files changed, 28 insertions(+), 34 deletions(-) diff --git a/tests/flytekit/unit/core/test_data.py b/tests/flytekit/unit/core/test_data.py index d686fae686..667445321b 100644 --- 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 index af2b2bf6f8..2fc8b6c452 100644 --- 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,7 +33,7 @@ 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", @@ -42,7 +41,6 @@ def test_initialise_azure_file_provider_with_service_principal(): "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)