""" Tests for chat completions router. """ import json import pytest from fastapi.testclient import TestClient from httpx import AsyncClient from src.chat import constants @pytest.mark.unit def test_chat_completion_non_streaming( client: TestClient, mock_chat_request: dict, ) -> None: """Test non-streaming chat completion.""" response = client.post("/v1/chat/completions", json=mock_chat_request) assert response.status_code == 200 data = response.json() # Verify response structure assert "id" in data assert data["object"] == constants.CHAT_COMPLETION_OBJECT assert "created" in data assert data["model"] == mock_chat_request["model"] assert len(data["choices"]) == 1 # Verify choice structure choice = data["choices"][0] assert choice["index"] == 0 assert choice["message"]["role"] == constants.ROLE_ASSISTANT assert choice["message"]["content"] # Should have content assert choice["finish_reason"] == constants.FINISH_REASON_STOP # Verify usage assert "usage" in data assert data["usage"]["prompt_tokens"] > 0 assert data["usage"]["completion_tokens"] > 0 assert data["usage"]["total_tokens"] > 0 @pytest.mark.unit def test_chat_completion_validation_error(client: TestClient) -> None: """Test chat completion with invalid request.""" # Missing required field 'messages' invalid_request = {"model": "Tatlock"} response = client.post("/v1/chat/completions", json=invalid_request) assert response.status_code == 422 data = response.json() assert "error" in data assert data["error"]["type"] == "invalid_request_error" @pytest.mark.unit async def test_chat_completion_streaming( async_client: AsyncClient, mock_streaming_chat_request: dict, ) -> None: """Test streaming chat completion with 20s timeout per turn.""" import asyncio async def read_stream_with_timeout(): """Read SSE stream with timeout protection.""" chunks = [] async with async_client.stream( "POST", "/v1/chat/completions", json=mock_streaming_chat_request, timeout=20.0, # 20 second timeout per turn ) as response: assert response.status_code == 200 assert response.headers["content-type"] == "text/event-stream; charset=utf-8" async for line in response.aiter_lines(): # Skip empty lines if not line.strip(): continue # Parse SSE format: "data: {json}" if line.startswith("data: "): data_str = line[6:].strip() # Remove "data: " prefix # Check for [DONE] marker if data_str == "[DONE]": break # Parse JSON chunk try: chunk = json.loads(data_str) chunks.append(chunk) # Verify chunk structure assert chunk["object"] == constants.CHAT_COMPLETION_CHUNK_OBJECT assert chunk["model"] == mock_streaming_chat_request["model"] assert len(chunk["choices"]) == 1 except json.JSONDecodeError as e: pytest.fail(f"Invalid JSON in stream: {data_str[:100]}... Error: {e}") return chunks # Execute with overall timeout try: chunks = await asyncio.wait_for(read_stream_with_timeout(), timeout=20.0) except TimeoutError: pytest.fail("Streaming test timed out after 20 seconds") # Verify we got chunks assert len(chunks) > 0, "No chunks received from stream" # First chunk should have role first_chunk = chunks[0] assert first_chunk["choices"][0]["delta"].get("role") == constants.ROLE_ASSISTANT # Last chunk should have finish_reason last_chunk = chunks[-1] assert last_chunk["choices"][0]["finish_reason"] == constants.FINISH_REASON_STOP @pytest.mark.unit def test_chat_completion_temperature_validation( client: TestClient, mock_chat_request: dict, ) -> None: """Test temperature parameter validation.""" # Invalid temperature (too high) invalid_request = {**mock_chat_request, "temperature": 3.0} response = client.post("/v1/chat/completions", json=invalid_request) assert response.status_code == 422 # Valid temperature valid_request = {**mock_chat_request, "temperature": 0.5} response = client.post("/v1/chat/completions", json=valid_request) assert response.status_code == 200 @pytest.mark.unit def test_chat_completion_message_roles( client: TestClient, mock_chat_request: dict, ) -> None: """Test different message roles.""" request_with_system = { **mock_chat_request, "messages": [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello!"}, ], } response = client.post("/v1/chat/completions", json=request_with_system) assert response.status_code == 200 @pytest.mark.unit def test_chat_completion_invalid_role( client: TestClient, mock_chat_request: dict, ) -> None: """Test invalid message role.""" invalid_request = { **mock_chat_request, "messages": [{"role": "invalid_role", "content": "test"}], } response = client.post("/v1/chat/completions", json=invalid_request) assert response.status_code == 422