Progress towards validation

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.
+        An `UNPROCESSABLE_ENTITY` response if the image is given and is not
+        either a PNG or a JPEG.
+    """
+    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)
+    file_size_bytes = len(pil_image.tobytes())
+
+    if pil_image.format != 'PNG':
+        return wrapped(*args, **kwargs)
+
+    documented_max_png_bytes = 2 * 1024 * 1024
+    if file_size_bytes > 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

@@ -1080,6 +1080,9 @@ def test_png(
         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