Skip to content

Commit

Permalink
SAASMLOPS-1058-2 Reformat code
Browse files Browse the repository at this point in the history
Signed-off-by: Alex Vinnik <[email protected]>
  • Loading branch information
alex-vinnik-sp committed Jan 22, 2024
1 parent 4fb6a79 commit 00be0e4
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 11 deletions.
1 change: 0 additions & 1 deletion sdk/python/feast/infra/offline_stores/snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
24 changes: 14 additions & 10 deletions sdk/python/tests/unit/infra/offline_stores/test_snowflake.py
Original file line number Diff line number Diff line change
@@ -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",
Expand Down Expand Up @@ -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://*"

0 comments on commit 00be0e4

Please sign in to comment.