diff --git a/tests/mp/io/loader/test_s3_loader.py b/tests/mp/io/loader/test_s3_loader.py index fa4161e..a6a0c92 100644 --- a/tests/mp/io/loader/test_s3_loader.py +++ b/tests/mp/io/loader/test_s3_loader.py @@ -13,18 +13,19 @@ def test_key_exists_true(mocker): - setup_key_exists(mocker) + setup_key_exists(mocker, ret_val={'ResponseMetadata': {'HTTPStatusCode': 200}}) assert s3_loader.key_exists(mock_event_keys[0]) def test_key_exists_false(mocker): - setup_key_exists(mocker) + setup_key_exists(mocker, ret_val={'ResponseMetadata': {'HTTPStatusCode': 404}}) assert not s3_loader.key_exists('/aaa/bbb/ccc') def test_key_exists_diff_region(mocker): - s3_lister = setup_key_exists(mocker) - assert s3_loader.key_exists(mock_event_keys[0], region=MOCK_REGION_NAME) + not_found = {'ResponseMetadata': {'HTTPStatusCode': 404}} + s3_lister = setup_key_exists(mocker, ret_val=not_found) + assert not s3_loader.key_exists(mock_event_keys[0], region=MOCK_REGION_NAME) boto3.client.assert_called_with('s3', region_name=MOCK_REGION_NAME) s3_lister.list_objects_v2.assert_called_with( Bucket=MOCK_BUCKET_NAME, MaxKeys=1, Prefix=mock_event_keys[0] @@ -56,10 +57,10 @@ def test_download_file_from_s3_not_found(mocker): s3_getter.download_file.assert_called_with(dest) -def setup_key_exists(mocker): +def setup_key_exists(mocker, ret_val={}): mocker.patch.dict(os.environ, {SOURCE_BUCKET: MOCK_BUCKET_NAME}) s3_lister = MagicMock() - s3_lister.list_objects_v2 = MagicMock(return_value=mock_event_keys) + s3_lister.list_objects_v2 = MagicMock(return_value=ret_val) mocker.patch.object(boto3, 'client', MagicMock(return_value=s3_lister)) return s3_lister