Merge pull request #666 from adamtheturtle/max-image-size-png

Add validation for PNG maximum validation size

Author: adamtheturtle

src/mock_vws/_mock_web_query_api.py

@@ -13,6 +13,7 @@
 from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import pytz
+import requests
 import wrapt
 from PIL import Image
 from requests import codes
@@ -69,6 +70,53 @@ def validate_project_state(
     )
 
 
+@wrapt.decorator
+def validate_image_file_size(
+    wrapped: Callable[..., str],
+    instance: Any,  # pylint: disable=unused-argument
+    args: Tuple[_RequestObjectProxy, _Context],
+    kwargs: Dict,
+) -> str:
+    """
+    Validate the file size of the image given to the query endpoint.
+
+    Args:
+        wrapped: An endpoint function for `requests_mock`.
+        instance: The class that the endpoint function is in.
+        args: The arguments given to the endpoint function.
+        kwargs: The keyword arguments given to the endpoint function.
+
+    Returns:
+        The result of calling the endpoint.
+
+    Raises:
+        requests.exceptions.ConnectionError: The image file size is too large.
+    """
+    request, _ = args
+    body_file = io.BytesIO(request.body)
+
+    _, pdict = cgi.parse_header(request.headers['Content-Type'])
+    parsed = cgi.parse_multipart(
+        fp=body_file,
+        pdict={
+            'boundary': pdict['boundary'].encode(),
+        },
+    )
+
+    [image] = parsed['image']
+
+    image_file = io.BytesIO(image)
+    pil_image = Image.open(image_file)
+
+    if pil_image.format != 'PNG':
+        return wrapped(*args, **kwargs)
+
+    documented_max_png_bytes = 2 * 1024 * 1024
+    if len(image) > documented_max_png_bytes:
+        raise requests.exceptions.ConnectionError
+    return wrapped(*args, **kwargs)
+
+
 @wrapt.decorator
 def validate_image_format(
     wrapped: Callable[..., str],
@@ -626,6 +674,7 @@ def decorator(method: Callable[..., str]) -> Callable[..., str]:
             validate_date_header_given,
             validate_include_target_data,
             validate_max_num_results,
+            validate_image_file_size,
             validate_image_file_contents,
             validate_image_format,
             validate_image_field_given,

tests/mock_vws/test_query.py

@@ -25,6 +25,7 @@
     add_target_to_vws,
     delete_target,
     get_vws_target,
+    make_image_file,
     query,
     update_target,
     wait_for_target_processed,
@@ -1065,11 +1066,70 @@ class TestMaximumImageSize:
     Tests for maximum image sizes.
     """
 
-    def test_png(self) -> None:
+    def test_png(
+        self,
+        vuforia_database_keys: VuforiaDatabaseKeys,
+    ) -> None:
         """
-        See https://github.com/adamtheturtle/vws-python/issues/357 for
-        implementing this test.
+        According to
+        https://library.vuforia.com/articles/Solution/How-To-Perform-an-Image-Recognition-Query.
+        the maximum file size is "2MiB for PNG".
+
+        Above this limit, a ``ConnectionError`` is raised.
+        We do not test exactly at this limit, but that may be beneficial in the
+        future.
         """
+        documented_max_bytes = 2 * 1024 * 1024
+        width = height = 835
+        png_not_too_large = make_image_file(
+            file_format='PNG',
+            color_space='RGB',
+            width=width,
+            height=height,
+        )
+
+        image_content = png_not_too_large.getvalue()
+        body = {'image': ('image.jpeg', image_content, 'image/jpeg')}
+
+        image_content_size = len(image_content)
+        # We check that the image we created is just slightly smaller than the
+        # maximum file size.
+        #
+        # This is just because of the implementation details of ``image_file``.
+        assert image_content_size < documented_max_bytes
+        assert (image_content_size * 1.05) > documented_max_bytes
+
+        response = query(
+            vuforia_database_keys=vuforia_database_keys,
+            body=body,
+        )
+
+        assert_query_success(response=response)
+        assert response.json()['results'] == []
+
+        width = height = 836
+        png_not_too_large = make_image_file(
+            file_format='PNG',
+            color_space='RGB',
+            width=width,
+            height=height,
+        )
+
+        image_content = png_not_too_large.getvalue()
+        body = {'image': ('image.jpeg', image_content, 'image/jpeg')}
+        image_content_size = len(image_content)
+        # We check that the image we created is just slightly larger than the
+        # maximum file size.
+        #
+        # This is just because of the implementation details of ``image_file``.
+        assert image_content_size > documented_max_bytes
+        assert (image_content_size * 0.95) < documented_max_bytes
+
+        with pytest.raises(requests.exceptions.ConnectionError):
+            query(
+                vuforia_database_keys=vuforia_database_keys,
+                body=body,
+            )
 
     def test_jpeg(self) -> None:
         """