Add test with complex payload

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.6.0.dev6"
+version = "0.6.0.dev7"
 description = ""
 authors = ["Guillaume Aquilina <guillaume@workflowai.com>"]
 readme = "README.md"

tests/integration/conftest.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any, Optional
+from collections.abc import Callable
+from typing import Any, Optional, Union
 from unittest.mock import patch
 
 import pytest
@@ -51,34 +52,70 @@ def mock_response(
             status_code=status_code,
         )
 
-    def mock_stream(self, task_id: str = "city-to-capital"):
+    def mock_stream(
+        self,
+        task_id: str = "city-to-capital",
+        outputs: Optional[list[dict[str, Any]]] = None,
+        run_id: str = "1",
+        metadata: Optional[dict[str, Any]] = None,
+    ):
+        outputs = outputs or [
+            {"capital": ""},
+            {"capital": "Tok"},
+            {"capital": "Tokyo"},
+        ]
+        if metadata is None:
+            metadata = {"cost_usd": 0.01, "duration_seconds": 10.1}
+
+        payloads = [{"id": run_id, "task_output": o} for o in outputs]
+
+        final_payload = {**payloads[-1], **metadata}
+        payloads.append(final_payload)
+        streams = [f"data: {json.dumps(p)}\n\n".encode() for p in payloads]
+
         self.httpx_mock.add_response(
             url=f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run",
-            stream=IteratorStream(
-                [
-                    b'data: {"id":"1","task_output":{"capital":""}}\n\n',
-                    b'data: {"id":"1","task_output":{"capital":"Tok"}}\n\ndata: {"id":"1","task_output":{"capital":"Tokyo"}}\n\n',  # noqa: E501
-                    b'data: {"id":"1","task_output":{"capital":"Tokyo"},"cost_usd":0.01,"duration_seconds":10.1}\n\n',
-                ],
-            ),
+            stream=IteratorStream(streams),
         )
 
+    def check_register(
+        self,
+        task_id: str = "city-to-capital",
+        input_schema: Optional[Union[dict[str, Any], Callable[[dict[str, Any]], None]]] = None,
+        output_schema: Optional[Union[dict[str, Any], Callable[[dict[str, Any]], None]]] = None,
+    ):
+        request = self.httpx_mock.get_request(url=self.REGISTER_URL)
+        assert request is not None
+        assert request.headers["Authorization"] == "Bearer test"
+        assert request.headers["Content-Type"] == "application/json"
+        assert request.headers["x-workflowai-source"] == "sdk"
+        assert request.headers["x-workflowai-language"] == "python"
+
+        body = json.loads(request.content)
+        assert body["id"] == task_id
+        if callable(input_schema):
+            input_schema(body["input_schema"])
+        else:
+            assert body["input_schema"] == input_schema or {"city": {"type": "string"}}
+        if callable(output_schema):
+            output_schema(body["output_schema"])
+        else:
+            assert body["output_schema"] == output_schema or {"capital": {"type": "string"}}
+
     def check_request(
         self,
         version: Any = "production",
         task_id: str = "city-to-capital",
         task_input: Optional[dict[str, Any]] = None,
         **matchers: Any,
     ):
+        if not matchers:
+            matchers = {"url": f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run"}
         request = self.httpx_mock.get_request(**matchers)
         assert request is not None
-        assert request.url == f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run"
         body = json.loads(request.content)
-        assert body == {
-            "task_input": task_input or {"city": "Hello"},
-            "version": version,
-            "stream": False,
-        }
+        assert body["task_input"] == task_input or {"city": "Hello"}
+        assert body["version"] == version
         assert request.headers["Authorization"] == "Bearer test"
         assert request.headers["Content-Type"] == "application/json"
         assert request.headers["x-workflowai-source"] == "sdk"

tests/integration/run_test.py

@@ -112,27 +112,23 @@ async def city_to_capital(task_input: CityToCapitalTaskInput) -> CityToCapitalTa
         "input_schema": {
             "properties": {
                 "city": {
-                    "title": "City",
                     "type": "string",
                 },
             },
             "required": [
                 "city",
             ],
-            "title": "CityToCapitalTaskInput",
             "type": "object",
         },
         "output_schema": {
             "properties": {
                 "capital": {
-                    "title": "Capital",
                     "type": "string",
                 },
             },
             "required": [
                 "capital",
             ],
-            "title": "CityToCapitalTaskOutput",
             "type": "object",
         },
     }