test: multi format images test

Author: anna-singleton-resolver

.gitignore

@@ -14,3 +14,6 @@ docs/_build/
 
 # Generated protobuf code
 src/resolver_athena_client/generated/
+
+# Generated functional test images
+tests/test_support/images/*

README.md

@@ -123,9 +123,19 @@ bash scripts/compile_proto.sh
 ```
 
 ### Functional Tests
-Functional tests require an Athena environment to run against. You can set up
-the environment variables in a `.env` file in the root of the repository, or in
-your shell environment:
+Functional tests require an Athena environment to run against.
+
+#### Pre-Requisites
+You will need:
+- An Athena host URL.
+- An OAuth client ID and secret with access to the Athena environment.
+- An affiliate with Athena enabled.
+- `imagemagick` installed on your system and on your path at `magick`.
+
+
+#### Preparing your environment
+You can set up the environment variables in a `.env` file in the root of the
+repository, or in your shell environment:
 
 You must set the following variables:
 ```

tests/functional/conftest.py

@@ -1,12 +1,33 @@
 import os
+import shutil
+import subprocess
 import uuid
+from pathlib import Path
 
 import pytest
 import pytest_asyncio
 from dotenv import load_dotenv
 
 from resolver_athena_client.client.athena_options import AthenaOptions
 from resolver_athena_client.client.channel import CredentialHelper
+from tests.utils.image_generation import create_test_image
+
+IMAGEMAGICK_FORMATS = [
+    "gif",
+    "bmp",
+    "dib",
+    "png",
+    "webp",
+    "pbm",
+    "pgm",
+    "ppm",
+    "pxm",
+    "pnm",
+    "sr",
+    "ras",
+    "tiff",
+    "pic",
+]
 
 
 def get_required_env_var(name: str) -> str:
@@ -67,3 +88,58 @@ def athena_options() -> AthenaOptions:
         keepalive_interval=30.0,  # Longer intervals for persistent streams
         affiliate="Crisp",
     )
+
+
+@pytest.fixture
+def formatted_images() -> list[tuple[bytes, str]]:
+    if (magick_path := shutil.which("magick")) is None:
+        msg = (
+            "ImageMagick 'magick' command not found - cannot run "
+            "multi-format test"
+        )
+        raise AssertionError(msg)
+
+    images = []
+    this_dir = Path(__file__).resolve()
+    image_dir = this_dir.parent / "../test_support/images/"
+
+    if not image_dir.exists():
+        image_dir.mkdir(parents=True)
+    if not image_dir.is_dir():
+        msg = f"Image directory {image_dir} is not a directory"
+        raise AssertionError(msg)
+
+    base_image_format = "png"
+    base_image = create_test_image(448, 448, img_format=base_image_format)
+    base_image_path = image_dir / "base_image.png"
+    with base_image_path.open("wb") as f:
+        f.write(base_image)
+
+    for img_format in IMAGEMAGICK_FORMATS:
+        if img_format == base_image_format:
+            continue  # base image is already generated.
+
+        converted_image_path = image_dir / f"test_image.{img_format}"
+        if converted_image_path.exists():
+            continue  # already generated - probably from previous run.
+
+        cmd = f'magick "{base_image_path}" "{converted_image_path}"'
+        subprocess.run(  # noqa: S603 - false positive :(
+            [magick_path, str(base_image_path), str(converted_image_path)],
+            check=True,
+            shell=False,
+        )
+
+        if not converted_image_path.exists():
+            msg = f"Failed to create {img_format} image with command: {cmd}"
+            raise AssertionError(msg)
+
+    for path in image_dir.iterdir():
+        if path.is_file():
+            with path.open("rb") as f:
+                img_bytes = f.read()
+                images.append((img_bytes, path.suffix.lstrip(".")))
+
+    raw_uint8 = create_test_image(448, 448, img_format="raw_uint8")
+    images.append((raw_uint8, "raw_uint8"))
+    return images

tests/functional/test_multi_format.py

@@ -0,0 +1,47 @@
+import pytest
+
+from resolver_athena_client.client.athena_client import AthenaClient
+from resolver_athena_client.client.athena_options import AthenaOptions
+from resolver_athena_client.client.channel import (
+    CredentialHelper,
+    create_channel_with_credentials,
+)
+from resolver_athena_client.client.models import ImageData
+
+
+@pytest.mark.asyncio
+@pytest.mark.functional
+async def test_classsify_single_multi_format(
+    athena_options: AthenaOptions,
+    credential_helper: CredentialHelper,
+    formatted_images: list[tuple[bytes, str]],
+) -> None:
+    # Create gRPC channel with credentials
+    channel = await create_channel_with_credentials(
+        athena_options.host, credential_helper
+    )
+
+    failed_formats: list[tuple[str, Exception]] = []
+
+    async with AthenaClient(channel, athena_options) as client:
+        for img_bytes, img_format in formatted_images:
+            try:
+                # Create a unique test image for each iteration
+                image_data = ImageData(img_bytes)
+
+                # Classify with auto-generated correlation ID
+                result = await client.classify_single(image_data)
+
+                if result.error.code:
+                    msg = f"Image Result Error: {result.error.message}"
+                    pytest.fail(msg)
+
+            except Exception as e:  # noqa: BLE001, PERF203 - its a test.
+                failed_formats.append((img_format, e))
+
+    if failed_formats:
+        failure_messages = ", ".join(
+            [f"{fmt}: {err}" for fmt, err in failed_formats]
+        )
+        msg = f"Classification failed for formats: {failure_messages}"
+        raise AssertionError(msg)

tests/utils/image_generation.py

@@ -25,7 +25,9 @@ def _get_cached_image(
     return _image_cache[key]
 
 
-def create_random_image(width: int = 160, height: int = 120) -> bytes:
+def create_random_image(
+    width: int = 160, height: int = 120, img_format: str = "PNG"
+) -> bytes:
     """Create a minimal random image optimized for maximum speed.
 
     Args:
@@ -55,9 +57,24 @@ def create_random_image(width: int = 160, height: int = 120) -> bytes:
     x2, y2 = (width * 3) // 4, (height * 3) // 4
     draw.rectangle([x1, y1, x2, y2], fill=accent_color)
 
+    if img_format.upper() == "RAW_UINT8":
+        # Return raw bytes
+        r_bytes = []
+        g_bytes = []
+        b_bytes = []
+        for y in range(height):
+            for x in range(width):
+                pixel = image.getpixel((x, y))
+                assert pixel is not None
+                assert isinstance(pixel, tuple)
+                r_bytes.append(pixel[0])
+                g_bytes.append(pixel[1])
+                b_bytes.append(pixel[2])
+        return bytes(r_bytes + g_bytes + b_bytes)
+
     # Convert to PNG bytes
     buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
+    image.save(buffer, format=img_format)
     return buffer.getvalue()
 
 
@@ -142,20 +159,25 @@ async def iter_images(
 
 
 def create_test_image(
-    width: int = 160, height: int = 120, seed: int | None = None
+    width: int = 160,
+    height: int = 120,
+    seed: int | None = None,
+    img_format: str = "PNG",
 ) -> bytes:
     """Create a test image with specified dimensions and optional seed.
 
     Args:
         width: Width of the test image in pixels (default: 160)
         height: Height of the test image in pixels (default: 120)
         seed: Optional seed for reproducible image generation
+        img_format: Image format (default: PNG). Other formats like JPEG are
+            also supported.
 
     Returns:
-        PNG image bytes
+        image bytes in specified format (default: PNG)
 
     """
     if seed is not None:
         _rng.seed(seed)
 
-    return create_random_image(width, height)
+    return create_random_image(width, height, img_format)