refactor(handler,-utests): updates, refactoring, testing

.gitignore

@@ -3,6 +3,7 @@ python
 build
 node_modules/**
 .vscode
+.idea
 .pytest_cache/
-
+.DS_Store
 __pycache__

imaginex_lambda/handler.py

@@ -1,9 +1,11 @@
+import logging
 import os
 import shutil
 from contextlib import ExitStack
 from tempfile import TemporaryFile
-from typing import IO
+from typing import IO, Tuple, Dict, Any
 from urllib.request import urlopen
+from io import BytesIO
 
 import botocore.session
 from PIL import Image
@@ -19,28 +21,45 @@
 session = botocore.session.get_session()
 s3_client = session.create_client('s3')
 
-
 # @TODO: Add placeholder image for errors.
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
 
-def download_image(buffer: IO[bytes], img_url: str):
-    print("Downloading image...")
+def download_image(buffer: IO[bytes], img_url: str) -> Tuple[IO[bytes], Dict[str, Any]]:
+    """
+    Function responsible for downloading an image file from a given URL and writing its contents to a buffer.
+    It takes two arguments, buffer and img_url, and returns a dictionary containing information about the downloaded
+    image.
+
+    buffer: IO[bytes]
+        A file-like object that the downloaded image will be written to.
+    img_url: str
+        A string representing the URL of the image to be downloaded.
+    """
+    logger.info("Downloading image from %s", img_url)
 
     with urlopen(img_url) as r:
         content_type = r.headers['content-type']
         content_size = int(r.headers['content-length'])
 
         shutil.copyfileobj(r, buffer, DOWNLOAD_CHUNK_SIZE)
 
-    print("Downloaded!")
-    return {'content_type': content_type, 'content_size': content_size}
+    logger.info("Downloaded image from %s. Content type: %s, content size: %d", img_url, content_type, content_size)
+    return buffer, {'content_type': content_type, 'content_size': content_size}
 
 
-def get_s3_image(buffer: IO[bytes], key: str):
+def get_s3_image(buffer: IO[bytes], key: str) -> Tuple[IO[bytes], Dict[str, Any]]:
+    """
+    Function responsible for downloading an image file from an Amazon S3 bucket and writing its contents to a buffer.
+    It takes two arguments, buffer and key, and returns a dictionary containing information about the downloaded image.
+    """
     if not S3_BUCKET_NAME:
         raise Exception('must specify a value for S3_BUCKET_NAME for S3 support')
 
-    print("Downloading image from S3...")
+    logger.info("Downloading image from S3 with key: %s", key)
     r = s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=key)
 
     content_type = r['ContentType']
@@ -49,27 +68,41 @@ def get_s3_image(buffer: IO[bytes], key: str):
     with r['Body'] as fin:
         shutil.copyfileobj(fin, buffer, DOWNLOAD_CHUNK_SIZE)
 
-    return {'content_type': content_type, 'content_size': content_size}
+    logger.info("Downloaded image from S3 with key: %s. Content type: %s, content size: %d", key, content_type,
+                content_size)
+    return buffer, {'content_type': content_type, 'content_size': content_size}
 
 
-def optimize_image(buffer: IO[bytes], ext: str, width: int, quality: int):
-    print("Optimizing image...")
+def optimize_image(buffer: IO[bytes], ext: str, width: int, quality: int) -> bytes:
+    """
+    The optimize_image function is designed to optimize an image that is passed in. It resizes the image
+    to the given width (if necessary), compresses the image to reduce its size, and returns the optimized image data.
+
+    buffer: IO[bytes]
+        buffer containing the image data to be optimized.
+    ext: str
+        the file extension of the image, which is used to specify the format when saving the optimized image.
+    width: int
+        the maximum width of the image. If the image is wider than this value, it will be resized to fit within
+        this width.
+    quality: int
+        the quality of the compressed image. A higher quality will result in a larger file size, while a lower quality
+        will result in a smaller file size.
+    """
+    logger.info("Optimizing image...")
     with ExitStack() as stack:
         img = stack.enter_context(Image.open(buffer))
         if width < img.width:
-            print("Resizing image...")
+            logger.info("Resizing image...")
             new_height = int(width * img.height / img.width)
-
-            print(width, new_height)
-
+            logger.info("New height: %d", new_height)
             img = stack.enter_context(img.resize((width, new_height)))
-            print("Resized!")
-
-        tmp = stack.enter_context(TemporaryFile())
+            logger.info("Resized image to width: %d and height: %d", width, new_height)
+        tmp = stack.enter_context(BytesIO())
         img.save(tmp, quality=quality, optimize=True, format=ext)
         tmp.seek(0)
 
-        print("Optimized!")
+        logger.info("Optimized image!")
         return tmp.read()
 
 
@@ -81,17 +114,17 @@ def handler(event, context):
     Any image processing logic should be performed by other functions, to make unit testing easier.
     """
     try:
-        print("Starting...")
+        logger.info("Lambda function started")
 
         qs = event['queryStringParameters']
         url = qs.get('url', None)
         width = int(qs.get('w', 0))
         quality = int(qs.get('q', 70))
 
-        print(url, width, quality)
+        logger.info(f"url={url}, width={width}, quality={quality}")
 
         image_data, content_type, optimization_ratio = download_and_optimize(url, quality, width)
-
+        logger.info("Returning success response")
         return success(image_data, {
             'Vary': 'Accept',
             'Content-Type': content_type,
@@ -103,7 +136,7 @@ def handler(event, context):
         return error(str(exc), code=500)
 
 
-def download_and_optimize(url: str, quality: int, width: int):
+def download_and_optimize(url: str, quality: int, width: int) -> Tuple[bytes, str, float]:
     """
     This is the function responsible for coordinating the download and optimization of the images. It should
     not concern itself with any lambda-specific information.
@@ -116,10 +149,10 @@ def download_and_optimize(url: str, quality: int, width: int):
 
     with TemporaryFile() as buffer:
         if is_absolute(url):
-            download_image(buffer, url)
+            buffer, _ = download_image(buffer, url)
         else:
             key = url.strip('/')
-            get_s3_image(buffer, key)
+            buffer, _ = get_s3_image(buffer, key)
 
         buffer.flush()
         original = os.stat(buffer.name).st_size
@@ -136,15 +169,20 @@ def download_and_optimize(url: str, quality: int, width: int):
 
     ratio = len(image_data) / original if original != 0 else 0
 
-    print("Returning data...")
+    logger.info("Returning image and metadata")
     return image_data, content_type, ratio
 
 
 if __name__ == '__main__':
-    print("Running test...")
+    from io import BytesIO
+    import base64
 
+    print("Running test...")
     context = {'queryStringParameters':
-               {'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'}}
-    handler(context, None)
+                   {'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'}}
+    res = handler(context, None)
+    img = res['body']
+    Image.open(BytesIO(base64.b64decode(img.encode()))).show()
+    print('end')

test/test_failures.py

@@ -2,7 +2,8 @@
 
 import pytest
 
-from imaginex_lambda.handler import handler
+from imaginex_lambda.handler import handler, download_and_optimize
+from imaginex_lambda.utils import HandlerError
 
 
 def test_no_url():
@@ -34,3 +35,25 @@ def test_s3_invalid_bucket():
         assert r['statusCode'] == 500
         assert r['headers']['Content-Type'] == 'application/json'
         assert "InvalidBucketName" in r['body']
+
+
+def test_download_and_optimize_with_invalid_url():
+    # arrange
+    url = ''
+    quality = 50
+    width = 100
+
+    # act and assert
+    with pytest.raises(HandlerError):
+        download_and_optimize(url, quality, width)
+
+
+def test_download_and_optimize_with_zero_width():
+    # arrange
+    url = 'https://example.com/image.png'
+    quality = 50
+    width = 0
+
+    # act and assert
+    with pytest.raises(HandlerError):
+        download_and_optimize(url, quality, width)

test/test_formats.py

@@ -0,0 +1,38 @@
+from io import BytesIO
+from tempfile import TemporaryFile
+from unittest.mock import patch
+from PIL import Image
+
+from imaginex_lambda.handler import download_and_optimize
+
+import pytest
+from unittest.mock import MagicMock
+
+
+@pytest.mark.parametrize('img_type,expected_type', [
+    ('PNG', 'image/png'),
+    ('JPEG', 'image/jpeg'),
+    # ('PPM', 'image/ppm'),
+    ('GIF', 'image/gif'),
+    ('TIFF', 'image/tiff'),
+    ('BMP', 'image/bmp'),
+])
+def test_download_and_optimize_formats(img_type, expected_type):
+    # arrange
+    original_w, original_h, new_q, new_w = 300, 300, 80, 200
+
+    sample_url = "https://example.com"
+    img = Image.new('RGB', (original_w, original_h), color=(255, 0, 0))
+    tmp_img = TemporaryFile()
+    img.save(tmp_img, format=img_type)
+    download_image_mock = MagicMock(return_value=(tmp_img, {}))
+
+    # act
+    with patch('imaginex_lambda.handler.download_image', download_image_mock):
+        opt_img_bytes, content_type, ratio = download_and_optimize(sample_url, new_q, new_w)
+    tmp_img.close()
+
+    # assert
+    opt_img = Image.open(BytesIO(opt_img_bytes))
+    assert content_type == expected_type
+    assert opt_img.width == new_w