diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 1ce482ff507..d3e84f12268 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -589,7 +589,6 @@ def to_remote_storage(self) -> List[str]: native_export_path = self.export_path.replace("s3gov://", "s3://") return self._get_file_names_from_copy_into(cursor, native_export_path) - def _get_file_names_from_copy_into(self, cursor, native_export_path) -> List[str]: file_name_column_index = [ idx for idx, rm in enumerate(cursor.description) if rm.name == "FILE_NAME" diff --git a/sdk/python/tests/unit/infra/offline_stores/test_snowflake.py b/sdk/python/tests/unit/infra/offline_stores/test_snowflake.py index 1b43ad9cf6a..afc3ae97aef 100644 --- a/sdk/python/tests/unit/infra/offline_stores/test_snowflake.py +++ b/sdk/python/tests/unit/infra/offline_stores/test_snowflake.py @@ -1,18 +1,17 @@ -from unittest.mock import Mock, MagicMock, patch, ANY -import pytest +import re +from unittest.mock import ANY, MagicMock, patch -from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig +import pytest from feast.infra.offline_stores.snowflake import ( SnowflakeOfflineStoreConfig, SnowflakeRetrievalJob, ) - +from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig from feast.repo_config import RepoConfig -import re -@pytest.fixture(params=['s3','s3gov']) +@pytest.fixture(params=["s3", "s3gov"]) def retrieval_job(request): offline_store_config = SnowflakeOfflineStoreConfig( type="snowflake.offline", @@ -44,10 +43,15 @@ def retrieval_job(request): def test_to_remote_storage(retrieval_job): stored_files = ["just a path", "maybe another"] - with patch.object(retrieval_job, "to_snowflake", return_value=None) as mock_to_snowflake, \ - patch.object(retrieval_job, "_get_file_names_from_copy_into", return_value=stored_files) as mock_get_file_names_from_copy: - assert retrieval_job.to_remote_storage() == stored_files, "should return the list of files" + with patch.object( + retrieval_job, "to_snowflake", return_value=None + ) as mock_to_snowflake, patch.object( + retrieval_job, "_get_file_names_from_copy_into", return_value=stored_files + ) as mock_get_file_names_from_copy: + assert ( + retrieval_job.to_remote_storage() == stored_files + ), "should return the list of files" mock_to_snowflake.assert_called_once() mock_get_file_names_from_copy.assert_called_once_with(ANY, ANY) native_path = mock_get_file_names_from_copy.call_args[0][1] - assert re.match(f"^s3://.*", native_path), "path should be s3://*" + assert re.match("^s3://.*", native_path), "path should be s3://*"