chore: update tests

ebridges · Sep 6, 2020 · d3bac63 · d3bac63
1 parent f83e40a
commit d3bac63
Showing 1 changed file with 7 additions and 6 deletions.
diff --git 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