From 2f11c7c203f79fcefd1185e89bc7be6af2d4ed82 Mon Sep 17 00:00:00 2001
From: Matthew Farrellee <matt@cs.wisc.edu>
Date: Thu, 30 Jan 2025 20:35:27 -0500
Subject: [PATCH] add test for user message w/ image.data content (#906)

# What does this PR do?

a test exists for image.url content, but not image.data content. this
adds the former.


## Test Plan

`LLAMA_STACK_BASE_URL=http://localhost:8321 pytest -v
tests/client-sdk/inference/test_inference.py`


## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the
other checks if that's the case).
- [x] Ran pre-commit to handle lint / formatting issues.
- [x] Read the [contributor
guideline](https://github.com/meta-llama/llama-stack/blob/main/CONTRIBUTING.md),
      Pull Request section?
- [ ] Updated relevant documentation.
- [x] Wrote necessary unit or integration tests.
---
 tests/client-sdk/inference/test_inference.py | 54 +++++++++++++-------
 1 file changed, 36 insertions(+), 18 deletions(-)

diff --git a/tests/client-sdk/inference/test_inference.py b/tests/client-sdk/inference/test_inference.py
index 6dff1be24..b10ede357 100644
--- a/tests/client-sdk/inference/test_inference.py
+++ b/tests/client-sdk/inference/test_inference.py
@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 import base64
-import os
+import pathlib
 
 import pytest
 from pydantic import BaseModel
@@ -57,13 +57,20 @@ def get_weather_tool_definition():
 
 
 @pytest.fixture
-def base64_image_url():
-    image_path = os.path.join(os.path.dirname(__file__), "dog.png")
-    with open(image_path, "rb") as image_file:
-        # Convert the image to base64
-        base64_string = base64.b64encode(image_file.read()).decode("utf-8")
-        base64_url = f"data:image/png;base64,{base64_string}"
-        return base64_url
+def image_path():
+    return pathlib.Path(__file__).parent / "dog.png"
+
+
+@pytest.fixture
+def base64_image_data(image_path):
+    # Convert the image to base64
+    return base64.b64encode(image_path.read_bytes()).decode("utf-8")
+
+
+@pytest.fixture
+def base64_image_url(base64_image_data, image_path):
+    # suffix includes the ., so we remove it
+    return f"data:image/{image_path.suffix[1:]};base64,{base64_image_data}"
 
 
 def test_text_completion_non_streaming(llama_stack_client, text_model_id):
@@ -371,20 +378,31 @@ def test_image_chat_completion_streaming(llama_stack_client, vision_model_id):
     assert any(expected in streamed_content for expected in {"dog", "puppy", "pup"})
 
 
-def test_image_chat_completion_base64_url(
-    llama_stack_client, vision_model_id, base64_image_url
+@pytest.mark.parametrize("type_", ["url", "data"])
+def test_image_chat_completion_base64(
+    llama_stack_client, vision_model_id, base64_image_data, base64_image_url, type_
 ):
+    image_spec = {
+        "url": {
+            "type": "image",
+            "image": {
+                "url": {
+                    "uri": base64_image_url,
+                },
+            },
+        },
+        "data": {
+            "type": "image",
+            "image": {
+                "data": base64_image_data,
+            },
+        },
+    }[type_]
+
     message = {
         "role": "user",
         "content": [
-            {
-                "type": "image",
-                "image": {
-                    "url": {
-                        "uri": base64_image_url,
-                    },
-                },
-            },
+            image_spec,
             {
                 "type": "text",
                 "text": "Describe what is in this image.",