tatlock/tests/chat/test_router.py

"""
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": "mistral-nemo:latest"}
    
    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 asyncio.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