From 3c1b02855a6a19e8e49809d49b99253a5c18b4fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A1s=20Schmelczer?= <andras@schmelczer.dev>
Date: Wed, 6 Apr 2022 17:21:18 +0200
Subject: [PATCH] Fix tests
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: András Schmelczer <andras@schmelczer.dev>
---
 good_ai/tests/open_s3/test_large_file.py | 52 +++++++++---------------
 1 file changed, 19 insertions(+), 33 deletions(-)

diff --git a/good_ai/tests/open_s3/test_large_file.py b/good_ai/tests/open_s3/test_large_file.py
index 3537f52..38a3fad 100644
--- a/good_ai/tests/open_s3/test_large_file.py
+++ b/good_ai/tests/open_s3/test_large_file.py
@@ -1,9 +1,9 @@
 import unittest
 from pathlib import Path
 from typing import Any
-from unittest.mock import Mock, create_autospec, patch
+from unittest.mock import Mock, patch
 
-import botocore.session
+import boto3
 
 PATH = Path(__file__).parent.resolve()
 
@@ -15,7 +15,7 @@ credentials = {
     "aws_access_key_id": "YOUR_ACCESS_KEY_ID",
     "aws_secret_access_key": "YOUR_VERY_SECRET_ACCESS_KEY",
     "large_files_bucket_name": "create_a_bucket_and_put_its_name_here",
-    "endpoint_url": "this is optional, for backblaze, use this: https://s3.us-west-002.backblazeb2.com",
+    "aws_endpoint_url": "this is optional, for backblaze, use this: https://s3.us-west-002.backblazeb2.com",
 }
 
 
@@ -29,13 +29,8 @@ class TestLargeFile(unittest.TestCase):
         self.assertRaises(ValueError, LargeFile, "test-file", "w+r")
         self.assertRaises(ValueError, LargeFile, "test-file", "test")
 
-    @patch("botocore.session")
-    def test_initialized_with_dict(self, session: Any) -> None:
-        session_mock = Mock()
-        session.get_session = create_autospec(
-            botocore.session.get_session, return_value=session_mock
-        )
-
+    @patch.object(boto3, "client")
+    def test_initialized_with_dict(self, client: Any) -> None:
         s3 = Mock()
         s3.list_objects_v2 = Mock(
             return_value={
@@ -55,25 +50,23 @@ class TestLargeFile(unittest.TestCase):
                 ]
             }
         )
-
-        session_mock.create_client = Mock(return_value=s3)
+        boto3.client = Mock(return_value=s3)
 
         LargeFile.configure_credentials(
             aws_region_name=credentials["aws_region_name"],
             aws_access_key_id=credentials["aws_access_key_id"],
             aws_secret_access_key=credentials["aws_secret_access_key"],
             large_files_bucket_name=credentials["large_files_bucket_name"],
-            endpoint_url=credentials["endpoint_url"],
+            aws_endpoint_url=credentials["aws_endpoint_url"],
         )
         lf = LargeFile("test-file")
-        session_mock.set_credentials.assert_called_once_with(
-            access_key=credentials["aws_access_key_id"],
-            secret_key=credentials["aws_secret_access_key"],
-        )
-        session_mock.create_client.assert_called_once_with(
+
+        boto3.client.assert_called_once_with(
             "s3",
+            aws_access_key_id=credentials["aws_access_key_id"],
+            aws_secret_access_key=credentials["aws_secret_access_key"],
             region_name=credentials["aws_region_name"],
-            endpoint_url=credentials["endpoint_url"],
+            endpoint_url=credentials["aws_endpoint_url"],
         )
 
         s3.list_objects_v2.assert_called_once_with(
@@ -84,13 +77,8 @@ class TestLargeFile(unittest.TestCase):
         self.assertEqual(lf._local_name, "test-file-2")
         self.assertEqual(lf._s3_name, "test-file/2")
 
-    @patch("botocore.session")
-    def test_initialized_with_file(self, session: Any) -> None:
-        session_mock = Mock()
-        session.get_session = create_autospec(
-            botocore.session.get_session, return_value=session_mock
-        )
-
+    @patch.object(boto3, "client")
+    def test_initialized_with_file(self, client: Any) -> None:
         s3 = Mock()
         s3.list_objects_v2 = Mock(
             return_value={
@@ -111,19 +99,17 @@ class TestLargeFile(unittest.TestCase):
             }
         )
 
-        session_mock.create_client = Mock(return_value=s3)
+        boto3.client = Mock(return_value=s3)
 
         LargeFile.configure_credentials_from_file(PATH / "../../example_secrets.ini")
         lf = LargeFile("test-file")
 
-        session_mock.set_credentials.assert_called_once_with(
-            access_key=credentials["aws_access_key_id"],
-            secret_key=credentials["aws_secret_access_key"],
-        )
-        session_mock.create_client.assert_called_once_with(
+        boto3.client.assert_called_once_with(
             "s3",
+            aws_access_key_id=credentials["aws_access_key_id"],
+            aws_secret_access_key=credentials["aws_secret_access_key"],
             region_name=credentials["aws_region_name"],
-            endpoint_url=credentials["endpoint_url"],
+            endpoint_url=credentials["aws_endpoint_url"],
         )
 
         assert s3.list_objects_v2.called