diff --git a/tests/test_output_structure.py b/tests/test_output_structure.py new file mode 100644 index 0000000..3d79f67 --- /dev/null +++ b/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)