Asynchronous reading of control file on s3, removing boto3.

livebridge/controldata/controlfile.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import aiobotocore
-import boto3
 import json
 import logging
 import os
@@ -91,7 +90,7 @@ async def load(self, path, *, resolve_auth=False):
         if not path.startswith("s3://"):
             body = self._load_from_file(path)
         else:
-            body = self._load_from_s3(path)
+            body = await self._load_from_s3(path)
         return yaml.load(body)
 
     def _load_from_file(self, path):
@@ -105,16 +104,15 @@ def _load_from_file(self, path):
 
         return body
 
-    def _load_from_s3(self, url):
+    async def _load_from_s3(self, url):
         bucket, key = url.split('/', 2)[-1].split('/', 1)
         logger.info("Loading control file from s3: {} - {}".format(bucket, key))
-        config = Config(signature_version="s3v4") if self.config["region"] in ["eu-central-1"] else None
-        client = boto3.client(
-            's3',
-            region_name=self.config["region"],
-            aws_access_key_id=self.config["access_key"] or None,
-            aws_secret_access_key=self.config["secret_key"] or None,
-            config=config,
-        )
-        control_file = client.get_object(Bucket=bucket, Key=key)
-        return control_file["Body"].read()
+        session = aiobotocore.get_session()
+        client = session.create_client('s3',
+             region_name=self.config["region"],
+             aws_secret_access_key=self.config["secret_key"] or None,
+             aws_access_key_id=self.config["access_key"] or None)
+        control_file = await client.get_object(Bucket=bucket, Key=key)
+        control_data = await control_file["Body"].read()
+        await client.close()
+        return control_data

requirements.txt

@@ -2,7 +2,6 @@ aiobotocore==0.3.3
 aiohttp==2.1.0
 asynctest==0.10.0
 bleach==2.0.0
-boto3==1.4.4
 pytest==3.1.2
 pytest-cov==2.5.1
 python-dateutil==2.6.0

tests/test_controlfile.py

@@ -112,25 +112,28 @@ async def test_load_from_file(self):
     async def test_load_by_s3_url(self):
         file_path = os.path.join(os.path.dirname(__file__), "files", "control.yaml")
         yaml_str = open(file_path).read()
-        self.control._load_from_s3 = lambda x: yaml_str
+        self.control._load_from_s3 = asynctest.CoroutineMock(return_value=yaml_str)
         control = await self.control.load("s3://foobaz/control.yaml", resolve_auth=True)
 
         assert control["auth"]["dev"]["api_key"] == "F00Baz"
         assert control["auth"]["live"]["api_key"] == "Foobar"
 
-    @asynctest.ignore_loop
-    def test_load_from_s3(self):
+    async def test_load_from_s3(self):
         s3_url = "s3://foobaz/control.yaml"
-        mock_client = unittest.mock.Mock()
-        mock_client.get_object.return_value = {"Body":StringIO("foo")}
-        with unittest.mock.patch("boto3.client") as patched:
-            patched.return_value = mock_client
-            control_data = self.control._load_from_s3(s3_url)
+        s3_content = asynctest.MagicMock()
+        s3_content.read = asynctest.CoroutineMock(return_value="foo")
+        mock_client = asynctest.MagicMock()
+        mock_client.get_object = asynctest.CoroutineMock(return_value={"Body": s3_content})
+        mock_client.close = asynctest.CoroutineMock(return_value=True)
+        mock_session = unittest.mock.MagicMock()
+        mock_session.create_client = unittest.mock.MagicMock(return_value=mock_client)
+        with asynctest.patch("aiobotocore.get_session") as patched:
+            patched.return_value = mock_session
+            control_data = await self.control._load_from_s3(s3_url)
             assert control_data == "foo"
             assert patched.call_count == 1
-            assert patched.call_args[0] == ('s3',)
-            assert patched.call_args[1]["region_name"] == "eu-central-1"
-            assert type(patched.call_args[1]["config"]) == botocore.config.Config
+            assert mock_session.create_client.call_args[0][0] == "s3"
+            assert mock_session.create_client.call_args[1]["region_name"] == "eu-central-1"
             mock_client.get_object.assert_called_once_with(Bucket='foobaz', Key='control.yaml')
 
     async def test_controlfile_without_auth(self):