From 89b7868e3446d04522b8699714295b9d9c157166 Mon Sep 17 00:00:00 2001
From: Matthew Farrellee <matt@cs.wisc.edu>
Date: Thu, 30 Jan 2025 15:44:24 -0500
Subject: [PATCH] merge test_image_chat_completion_base64_{url,data}

---
 tests/client-sdk/inference/test_inference.py | 58 +++++++-------------
 1 file changed, 21 insertions(+), 37 deletions(-)

diff --git a/tests/client-sdk/inference/test_inference.py b/tests/client-sdk/inference/test_inference.py
index 81cb0adf3..b10ede357 100644
--- a/tests/client-sdk/inference/test_inference.py
+++ b/tests/client-sdk/inference/test_inference.py
@@ -378,47 +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_data(
-    llama_stack_client, vision_model_id, base64_image_data
+@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": {
-                    "data": base64_image_data,
-                },
-            },
-            {
-                "type": "text",
-                "text": "Describe what is in this image.",
-            },
-        ],
-    }
-    response = llama_stack_client.inference.chat_completion(
-        model_id=vision_model_id,
-        messages=[message],
-        stream=False,
-    )
-    message_content = response.completion_message.content.lower().strip()
-    assert len(message_content) > 0
-
-
-def test_image_chat_completion_base64_url(
-    llama_stack_client, vision_model_id, base64_image_url
-):
-    message = {
-        "role": "user",
-        "content": [
-            {
-                "type": "image",
-                "image": {
-                    "url": {
-                        "uri": base64_image_url,
-                    },
-                },
-            },
+            image_spec,
             {
                 "type": "text",
                 "text": "Describe what is in this image.",