Fix tests added from rebase broken by QueryResult type change

Author: nv-alicheng

src/inference_endpoint/core/types.py

@@ -137,6 +137,10 @@ def __post_init__(self):
             msgspec.structs.force_setattr(
                 self, "response_output", tuple(self.response_output)
             )
+        elif isinstance(self.response_output, dict):
+            for k, v in self.response_output.items():
+                if isinstance(v, list):
+                    self.response_output[k] = tuple(v)
 
 
 class StreamChunk(msgspec.Struct, tag="stream_chunk", kw_only=True):

tests/integration/endpoint_client/test_http_client_streaming.py

@@ -130,9 +130,10 @@ async def test_streaming_with_futures(self, mock_http_echo_server, tmp_path):
                     assert result.response_output == f"Non-streaming request {idx}"
                 else:
                     assert result.id == f"stream-{idx}"
+                    assert list(result.response_output.keys()) == ["output"]
                     assert (
-                        "Streaming" in result.response_output
-                        or result.response_output == f"Streaming request {idx}"
+                        "".join(result.response_output["output"])
+                        == f"Streaming request {idx}"
                     )
 
         finally:
@@ -194,7 +195,10 @@ async def test_mixed_streaming_non_streaming(self, futures_http_client):
                 ]
             else:
                 # Streaming response
-                assert "".join(result.response_output) == "Streaming response test"
+                assert (
+                    "".join(result.response_output["output"])
+                    == "Streaming response test"
+                )
 
     @pytest.mark.asyncio
     async def test_concurrent_streaming_requests(self, futures_http_client):
@@ -224,7 +228,7 @@ async def test_concurrent_streaming_requests(self, futures_http_client):
         for idx, result in results:
             assert result.id == f"concurrent-stream-{idx}"
             assert (
-                "".join(result.response_output)
+                "".join(result.response_output["output"])
                 == f"Concurrent streaming request number {idx}"
             )
             assert result.error is None
@@ -265,7 +269,7 @@ async def track_resolution(future):
         assert resolution_count == 1
         assert resolved_result is not None
         assert resolved_result.id == "test-single-resolution"
-        assert resolved_result.response_output == (
+        assert resolved_result.response_output["output"] == (
             "Test",
             " single future resolution with multiple words",
         )
@@ -303,8 +307,10 @@ async def test_streaming_future_first_chunk_access(self, futures_http_client):
 
         # Get complete response
         result = await future
-        assert result.response_output[0] == "Test"
-        assert result.response_output[1] == " first chunk access functionality"
+        assert result.response_output["output"][0] == "Test"
+        assert (
+            result.response_output["output"][1] == " first chunk access functionality"
+        )
 
         # Test 2: Check if first chunk is ready after awaiting
         query2 = Query(
@@ -347,8 +353,7 @@ async def test_streaming_single_chunk_complete(self, futures_http_client):
 
         # Complete response should be the same
         result = await future
-        assert result.response_output[0] == "Hi"
-        assert result.response_output[1] == ""
+        assert result.response_output["output"] == ("Hi",)
 
     @pytest.mark.asyncio
     async def test_streaming_race_first_chunks(self, futures_http_client):
@@ -391,8 +396,8 @@ async def test_streaming_race_first_chunks(self, futures_http_client):
         # Get first complete
         winner = done.pop().result()
         assert winner.id.startswith("race-")
-        assert winner.response_output[0] == "Query"
-        assert winner.response_output[1].startswith(" number")
+        assert winner.response_output["output"][0] == "Query"
+        assert winner.response_output["output"][1].startswith(" number")
 
         # Clean up remaining
         for f in pending:
@@ -498,4 +503,4 @@ async def test_streaming_concurrent_mixed_lengths(self, futures_http_client):
                 assert first == prompt.split()[0] if prompt else ""
 
             result = await future
-            assert "".join(result.response_output) == prompt
+            assert "".join(result.response_output["output"]) == prompt

tests/integration/endpoint_client/test_worker_core.py

@@ -210,7 +210,7 @@ async def test_worker_streaming_request(
             assert responses[0].id == "test-streaming"
 
             # Verify final response
-            assert final_content == ("Stream", " this response please")
+            assert final_content == {"output": ("Stream", " this response please")}
 
             # Shutdown
             worker._shutdown = True
@@ -435,7 +435,9 @@ async def test_worker_streaming_first_last_token_metadata(
             # Verify final result metadata
             assert final_result.metadata.get("final_chunk") is True
             assert final_result.id == "test-first-last-token"
-            assert final_result.response_output == ("Hello", " world this is a test")
+            assert final_result.response_output == {
+                "output": ("Hello", " world this is a test")
+            }
 
             # Verify intermediate chunks (if any) don't have first_chunk set to True
             for chunk in stream_chunks[1:]:
@@ -520,7 +522,7 @@ async def test_worker_streaming_empty_chunks(
             final_response = responses[-1]
             assert final_response.metadata.get("final_chunk") is True
             assert final_response.id == "test-empty-chunks"
-            assert final_response.response_output == ("",)  # Empty content
+            assert final_response.response_output["output"] == ()  # Empty content
 
             # Shutdown
             worker._shutdown = True
@@ -884,7 +886,7 @@ async def test_worker_streaming_chunk_before_final_single_word(
             assert not isinstance(
                 last_response, StreamChunk
             ), "Last response should not be a StreamChunk"
-            assert last_response.response_output == ("Hi", "")
+            assert last_response.response_output == {"output": ("Hi",)}
             assert last_response.metadata.get("final_chunk") is True
             assert last_response.id == "test-single-word-order"
 
@@ -947,7 +949,7 @@ async def test_worker_streaming_chunk_before_final_single_word(
             assert not isinstance(
                 last_empty, StreamChunk
             ), "Last response should not be a StreamChunk"
-            assert last_empty.response_output == ("",)
+            assert last_empty.response_output == {"output": ()}
             assert last_empty.id == "test-empty-order"
 
             # Shutdown

tests/integration/endpoint_client/test_worker_errors.py

@@ -693,7 +693,7 @@ async def malformed_json_handler(request):
 
                 # Malformed JSON is skipped, so we get an empty response, not an error
                 assert response.error is None
-                assert response.response_output == ()
+                assert response.response_output == {"output": ()}
 
                 # Shutdown
                 worker._shutdown = True