♻️ refactor(handler): isolates the handler code from the image proces…

…sing

This was made to make unit testing and debugging easier, since the handler will swallow any
exception it catches.

imaginex_lambda/handler.py

@@ -8,7 +8,7 @@
 import botocore.session
 from PIL import Image
 
-from imaginex_lambda.utils import error, success, is_absolute, get_extension
+from imaginex_lambda.utils import error, success, is_absolute, get_extension, HandlerError
 
 DOWNLOAD_CHUNK_SIZE = int(os.getenv('DOWNLOAD_CHUNK_SIZE', 1024))
 S3_BUCKET_NAME = os.getenv('S3_BUCKET_NAME', None)
@@ -74,6 +74,12 @@ def optimize_image(buffer: IO[bytes], ext: str, width: int, quality: int):
 
 
 def handler(event, context):
+    """
+    Lambda function handler.
+
+    Its sole responsibility is to to parse the context and generate a formatted response, including error responses.
+    Any image processing logic should be performed by other functions, to make unit testing easier.
+    """
     try:
         print("Starting...")
 
@@ -84,40 +90,53 @@ def handler(event, context):
 
         print(url, width, quality)
 
-        if not url:
-            return error('url is required')
-        if width <= 0:
-            return error('width must be greater than zero')
-
-        with TemporaryFile() as buffer:
-            if is_absolute(url):
-                download_image(buffer, url)
-            else:
-                key = url.strip('/')
-                get_s3_image(buffer, key)
-
-            original = os.stat(buffer.name).st_size
-            mime = get_extension(buffer)
-            content_type = mime['content_type']
-            extension = mime['extension']
-
-            image_data = optimize_image(
-                buffer,
-                ext=extension,
-                width=width,
-                quality=quality
-            )
-
-        print("Returning data...")
+        image_data, content_type, optimization_ratio = download_and_optimize(url, quality, width)
+
         return success(image_data, {
             'Vary': 'Accept',
             'Content-Type': content_type,
-            'X-Optimization-Ratio': f'{len(image_data) / original:.4f}',
+            'X-Optimization-Ratio': f'{optimization_ratio:.4f}',
         })
+    except HandlerError as exc:
+        return error(str(exc), code=exc.code)
     except Exception as exc:
         return error(str(exc), code=500)
 
 
+def download_and_optimize(url: str, quality: int, width: int):
+    """
+    This is the function responsible for coordinating the download and optimization of the images. It should
+    not concern itself with any lambda-specific information.
+    """
+
+    if not url:
+        raise HandlerError('url is required')
+    if width <= 0:
+        raise HandlerError('width must be greater than zero')
+
+    with TemporaryFile() as buffer:
+        if is_absolute(url):
+            download_image(buffer, url)
+        else:
+            key = url.strip('/')
+            get_s3_image(buffer, key)
+
+        original = os.stat(buffer.name).st_size
+        mime = get_extension(buffer)
+        content_type = mime['content_type']
+        extension = mime['extension']
+
+        image_data = optimize_image(
+            buffer,
+            ext=extension,
+            width=width,
+            quality=quality
+        )
+
+    print("Returning data...")
+    return image_data, content_type, len(image_data) / original
+
+
 if __name__ == '__main__':
     print("Running test...")
 

imaginex_lambda/utils.py

@@ -6,6 +6,14 @@
 import filetype
 
 
+class HandlerError(Exception):
+    code: int
+
+    def __init__(self, msg: str, code=422) -> None:
+        super().__init__(msg)
+        self.code = code
+
+
 def success(image_data, headers):
     return {
         'statusCode': 200,

test/test_success.py

@@ -2,24 +2,43 @@
 
 import pytest
 
-from imaginex_lambda.handler import handler, S3_BUCKET_NAME
+from imaginex_lambda.handler import handler, S3_BUCKET_NAME, download_and_optimize
 
 
-@pytest.mark.parametrize('ratio,type,qs', [
-    ('0.0037', 'image/png',
-     {"q": "80", "w": "100", "url": "example.png"}),
-    ('0.0110', 'image/jpeg',
-     {"q": "40", "w": "250", "url": "https://s3.eu-central-1.amazonaws.com/fllite-dev-main/business_case_custom_images"
-                                    "/sun_valley_2_5f84953fef8c6_63a2668275433.jfif"}),
-    ('0.2109', 'image/jpeg',
-     {'q': "80", 'w': "100", 'url': "http://site.meishij.net/r/58/25/3568808/a3568808_142682562777944.jpg"}),
-])
-def test_handler_success(ratio, type, qs):
-    if not qs['url'].startswith('http') and S3_BUCKET_NAME is None:
-        pytest.skip('specify a value for S3_BUCKET_NAME to run S3 tests')
+def test_handler_success():
+    """
+    This test will mock the result of the 'download_and_optimize' function
+    to return a fixed value. This way we can check if the handler single responsibility
+    (to parse the context and generate a formatted response) is working correctly.
+
+    Specific image processing tests should call 'download_and_optimize' directly,
+    so it is easier to check for exceptions and edge cases.
+    """
+
+    context = {'queryStringParameters': {'url': 'abc.png', 'q': '50', 'w': '100'}}
+    fake_optimization_return = (b'abcdef', 'application/someimage', 0.3)
+
+    with patch('imaginex_lambda.handler.download_and_optimize') as p:
+        p.return_value = fake_optimization_return
+        r = handler(context, None)
 
-    r = handler({'queryStringParameters': qs}, None)
     assert r['statusCode'] == 200
     assert r['isBase64Encoded'] is True
-    assert r['headers']['Content-Type'] == type
-    assert r['headers']['X-Optimization-Ratio'] == ratio
+    assert r['headers']['Content-Type'] == 'application/someimage'
+    assert r['headers']['X-Optimization-Ratio'] == '0.3000'
+
+
+@pytest.mark.parametrize('expected_ratio,expected_type,q,w,url', [
+    (0.0037, 'image/png', 80, 100, "example.png"),
+    (0.0110, 'image/jpeg', 40, 250, "https://s3.eu-central-1.amazonaws.com/fllite-dev-main/"
+                                    "business_case_custom_images/sun_valley_2_5f84953fef8c6_63a2668275433.jfif"),
+    (0.2109, 'image/jpeg', 80, 100, "http://site.meishij.net/r/58/25/3568808/a3568808_142682562777944.jpg"),
+])
+def test_process_success(expected_ratio, expected_type, q, w, url):
+    if not url.startswith('http') and S3_BUCKET_NAME is None:
+        pytest.skip('specify a value for S3_BUCKET_NAME to run S3 tests')
+
+    image_data, content_type, ratio = download_and_optimize(url=url, quality=q, width=w)
+    assert isinstance(image_data, bytes)
+    assert content_type == expected_type
+    assert round(ratio, 4) == expected_ratio