Add custom load() for s3.Bucket

boto3/s3/inject.py

@@ -13,12 +13,40 @@
 from boto3.s3.transfer import S3Transfer
 from boto3 import utils
 
+from botocore.exceptions import ClientError
 
 def inject_s3_transfer_methods(class_attributes, **kwargs):
     utils.inject_attribute(class_attributes, 'upload_file', upload_file)
     utils.inject_attribute(class_attributes, 'download_file', download_file)
 
 
+def inject_bucket_load(class_attributes, **kwargs):
+    utils.inject_attribute(class_attributes, 'load', bucket_load)
+
+
+def bucket_load(self, *args, **kwargs):
+    """Fake s3.Bucket.load method.
+
+    This emulates a bucket load method such that you can use::
+
+        bucket = s3.Bucket('foo')
+        bucket.load()
+        # Or even just:
+        bucket.creation_date
+
+    """
+    # We can't actually get the bucket's attributes from a HeadBucket,
+    # so we need to use a ListBuckets and search for our bucket.
+    response = self.meta.client.list_buckets()
+    for bucket_data in response['Buckets']:
+        if bucket_data['Name'] == self.name:
+            self.meta.data = bucket_data
+            break
+    else:
+        raise ClientError({'Error': {'Code': '404', 'Message': 'NotFound'}},
+                          'ListBuckets')
+
+
 def upload_file(self, Filename, Bucket, Key, ExtraArgs=None,
                 Callback=None, Config=None):
     transfer = S3Transfer(self, Config)

boto3/session.py

@@ -279,6 +279,10 @@ def _register_default_handlers(self):
             'creating-client-class.s3',
             boto3.utils.lazy_call(
                 'boto3.s3.inject.inject_s3_transfer_methods'))
+        self._session.register(
+            'creating-resource-class.s3.Bucket',
+            boto3.utils.lazy_call(
+                'boto3.s3.inject.inject_bucket_load'))
         self._session.register(
             'creating-resource-class.dynamodb',
             boto3.utils.lazy_call(

tests/integration/test_s3.py

@@ -19,6 +19,7 @@
 import shutil
 import hashlib
 import string
+import datetime
 
 from tests import unittest, unique_id
 from botocore.compat import six
@@ -460,11 +461,41 @@ def test_transfer_methods_through_client(self):
         assert_files_equal(filename, download_path)
 
 
-class TestS3TransferMethodInjection(unittest.TestCase):
+class TestS3MethodInjection(unittest.TestCase):
     def test_transfer_methods_injected_to_client(self):
         session = boto3.session.Session(region_name='us-west-2')
         client = session.client('s3')
         self.assertTrue(hasattr(client, 'upload_file'),
                         'upload_file was not injected onto S3 client')
         self.assertTrue(hasattr(client, 'download_file'),
                         'download_file was not injected onto S3 client')
+
+    def test_bucket_resource_has_load_method(self):
+        session = boto3.session.Session(region_name='us-west-2')
+        bucket = session.resource('s3').Bucket('fakebucket')
+        self.assertTrue(hasattr(bucket, 'load'),
+                        'load() was not injected onto S3 Bucket resource.')
+
+
+class TestCustomS3BucketLoad(unittest.TestCase):
+    def setUp(self):
+        self.region = 'us-west-2'
+        self.session = boto3.session.Session(region_name=self.region)
+        self.s3 = self.session.resource('s3')
+        self.bucket_name = unique_id('boto3-test')
+
+    def create_bucket_resource(self, bucket_name, region=None):
+        if region is None:
+            region = self.region
+        kwargs = {'Bucket': bucket_name}
+        if region != 'us-east-1':
+            kwargs['CreateBucketConfiguration'] = {
+                'LocationConstraint': region
+            }
+        bucket = self.s3.create_bucket(**kwargs)
+        self.addCleanup(bucket.delete)
+        return bucket
+
+    def test_can_access_buckets_creation_date(self):
+        bucket = self.create_bucket_resource(random_bucket_name())
+        self.assertIsInstance(bucket.creation_date, datetime.datetime)

tests/unit/s3/test_inject.py

@@ -13,10 +13,12 @@
 from tests import unittest
 import mock
 
+from botocore.exceptions import ClientError
+
 from boto3.s3 import inject
 
 
-class TestInject(unittest.TestCase):
+class TestInjectTransferMethods(unittest.TestCase):
     def test_inject_upload_download_file_to_client(self):
         class_attributes = {}
         inject.inject_s3_transfer_methods(class_attributes=class_attributes)
@@ -41,3 +43,35 @@ def test_download_file_proxies_to_transfer_object(self):
             transfer.return_value.download_file.assert_called_with(
                 bucket='bucket', key='key', filename='filename',
                 extra_args=None, callback=None)
+
+
+class TestBucketLoad(unittest.TestCase):
+    def setUp(self):
+        self.client = mock.Mock()
+        self.resource = mock.Mock()
+        self.resource.meta.client = self.client
+
+    def test_bucket_load_finds_bucket(self):
+        self.resource.name = 'MyBucket'
+        self.client.list_buckets.return_value = {
+            'Buckets': [
+                {'Name': 'NotMyBucket', 'CreationDate': 1},
+                {'Name': self.resource.name, 'CreationDate': 2},
+            ],
+        }
+
+        inject.bucket_load(self.resource)
+        self.assertEqual(
+            self.resource.meta.data,
+            {'Name': self.resource.name, 'CreationDate': 2})
+
+    def test_bucket_load_raise_error(self):
+        self.resource.name = 'MyBucket'
+        self.client.list_buckets.return_value = {
+            'Buckets': [
+                {'Name': 'NotMyBucket', 'CreationDate': 1},
+                {'Name': 'NotMine2', 'CreationDate': 2},
+            ],
+        }
+        with self.assertRaises(ClientError):
+            inject.bucket_load(self.resource)