From bd19d08272bc5cc16d7474c3e3a5578a9c594d32 Mon Sep 17 00:00:00 2001
From: Bo Bayles <bbayles@gmail.com>
Date: Wed, 30 Oct 2024 11:31:40 -0500
Subject: [PATCH] Add S3 head_object helper

---
 boto3_helpers/s3.py | 37 ++++++++++++++++++++++++++++++++++
 tests/test_s3.py    | 49 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/boto3_helpers/s3.py b/boto3_helpers/s3.py
index 10fbf79..a423f33 100644
--- a/boto3_helpers/s3.py
+++ b/boto3_helpers/s3.py
@@ -82,3 +82,40 @@ def query_object(bucket, key, query, input_format, *, s3_client=None, **kwargs):
                     yield loads(line)
                 else:
                     data[:] = line
+
+
+def head_bucket(bucket, s3_client=None, **kwargs):
+    """Perform a ``HeadObject`` API call and return the response. If the given
+    *bucket* does not exist, raise ``s3_client.exceptions.NoSuchBucket``
+
+    * *bucket* is the S3 bucket to use
+    * *s3_client* is a ``boto3.client('s3')`` instance. If not given,
+      one will be created with ``boto3.client('s3')``.
+    * *kwargs* are passed to the ``head_object`` method.
+
+    The ``boto3`` docs infamously claim that the `head_object` method can raise
+    ``S3.Client.exceptions.NoSuchBucket``, leading reasonable people to assume that
+    this exception will be raised if the requested *bucket* does not exist. Alas,
+    it raises a generic ``ClientError`` instead. This function fixes the problem.
+
+    .. code-block:: python
+
+        from boto3 import client as boto3_client
+        from boto3_helpers.s3 import head_bucket
+
+        s3_client = boto3_client('s3')
+        try:
+            head_bucket('ExampleBucket', s3_client=s3_client)
+        except s3_client.exceptions.NoSuchBucket:
+            print('No such bucket')
+        else:
+            print('That bucket exists')
+    """
+    s3_client = s3_client or boto3_client('s3')
+    kwargs['Bucket'] = bucket
+    try:
+        return s3_client.head_bucket(**kwargs)
+    except s3_client.exceptions.ClientError as e:
+        if e.response['Error']['Code'] == '404':
+            raise s3_client.exceptions.NoSuchBucket(e.response, e.operation_name)
+        raise
diff --git a/tests/test_s3.py b/tests/test_s3.py
index 7874d14..1047f33 100644
--- a/tests/test_s3.py
+++ b/tests/test_s3.py
@@ -1,7 +1,10 @@
 from unittest import TestCase
 from unittest.mock import MagicMock
 
-from boto3_helpers.s3 import query_object
+from boto3 import client as boto3_client
+from botocore.stub import Stubber
+
+from boto3_helpers.s3 import head_bucket, query_object
 
 
 class QueryObjectTests(TestCase):
@@ -68,3 +71,47 @@ def test_custom(self):
             InputSerialization=input_serialization,
             OutputSerialization={'JSON': {}},
         )
+
+
+class HeadBucketTest(TestCase):
+    def test_exists(self):
+        mock_s3_client = boto3_client('s3', region_name='not-a-region')
+        stubber = Stubber(mock_s3_client)
+        params = {'Bucket': 'example'}
+        resp = {}
+        stubber.add_response('head_bucket', resp, params)
+
+        with stubber:
+            actual = head_bucket('example', s3_client=mock_s3_client)
+
+        self.assertEqual(actual, resp)
+
+    def test_not_exists(self):
+        mock_s3_client = boto3_client('s3', region_name='not-a-region')
+        stubber = Stubber(mock_s3_client)
+        params = {'Bucket': 'example'}
+        stubber.add_client_error(
+            'head_bucket',
+            service_error_code='404',
+            service_message='The specified bucket does not exist.',
+            http_status_code=404,
+            expected_params=params,
+        )
+
+        with stubber, self.assertRaises(mock_s3_client.exceptions.NoSuchBucket):
+            head_bucket('example', s3_client=mock_s3_client)
+
+    def test_other_error(self):
+        mock_s3_client = boto3_client('s3', region_name='not-a-region')
+        stubber = Stubber(mock_s3_client)
+        params = {'Bucket': 'example'}
+        stubber.add_client_error(
+            'head_bucket',
+            service_error_code='403',
+            service_message='Some other bad thing happened.',
+            http_status_code=403,
+            expected_params=params,
+        )
+
+        with stubber, self.assertRaises(mock_s3_client.exceptions.ClientError):
+            head_bucket('example', s3_client=mock_s3_client)