Add tests for verify output structure.

tests/test_output_structure.py

@@ -0,0 +1,74 @@
+import pytest
+import requests
+from typing import Dict, Any
+from .conftest import BASE_URL
+
+
+class TestOutputStructure:
+    """Response structure tests"""
+
+    def test_general_output_structure(self, api_client):
+        """Verify the output structure /general"""
+        payload = {
+            "text": "Apple Inc. in California",
+            "entities": ["organization", "location"],
+            "threshold": 0.3
+        }
+        response = api_client.post(f"{BASE_URL}/general", json=payload)
+        data = response.json()
+
+        if data["outputs"]:
+            entity = data["outputs"][0]
+            assert "entity" in entity
+            assert "span" in entity
+            assert "start" in entity
+            assert "end" in entity
+            assert "score" in entity
+            assert isinstance(entity["start"], int)
+            assert isinstance(entity["end"], int)
+            assert isinstance(entity["score"], float)
+
+    def test_relation_output_structure(self, api_client):
+        """Verify the output structure /relation-extraction"""
+        payload = {
+            "text": "Steve Jobs founded Apple",
+            "relations": ["founded"],
+            "threshold": 0.3
+        }
+        response = api_client.post(f"{BASE_URL}/relation-extraction", json=payload)
+        data = response.json()
+
+        if data["outputs"]:
+            relation = data["outputs"][0]
+            assert "source" in relation
+            assert "relation" in relation
+            assert "target" in relation
+            assert "score" in relation
+
+            # Verify structure source and target
+            for entity_key in ["source", "target"]:
+                entity = relation[entity_key]
+                assert "entity" in entity
+                assert "span" in entity
+                assert "start" in entity
+                assert "end" in entity
+                assert "score" in entity
+
+    def test_summarization_output_structure(self, api_client):
+        """Verify the output structure /summarization"""
+        payload = {
+            "text": "AI is changing the world with machine learning.",
+            "threshold": 0.3
+        }
+        response = api_client.post(f"{BASE_URL}/summarization", json=payload)
+        data = response.json()
+
+        if data["outputs"]:
+            summary = data["outputs"][0]
+            assert "text" in summary
+            assert "start" in summary
+            assert "end" in summary
+            assert "score" in summary
+            assert isinstance(summary["start"], int)
+            assert isinstance(summary["end"], int)
+            assert isinstance(summary["score"], float)