zzrouter/backend/tests/conftest.py

"""Test configuration and shared fixtures."""
import asyncio
import pytest
import aiosqlite
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock
from contextlib import asynccontextmanager

# Test database path
TEST_DB_PATH = Path(__file__).parent / "data" / "test.db"


@asynccontextmanager
async def get_test_db():
    """Get in-memory test database connection."""
    conn = await aiosqlite.connect(":memory:")
    conn.row_factory = aiosqlite.Row
    await _init_test_schema(conn)
    await _seed_test_data(conn)
    yield conn
    await conn.close()


async def _init_test_schema(db: aiosqlite.Connection):
    """Initialize test database schema."""
    await db.executescript("""
        CREATE TABLE IF NOT EXISTS client_keys (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key TEXT UNIQUE NOT NULL,
            name TEXT,
            is_active BOOLEAN DEFAULT TRUE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );

        CREATE TABLE IF NOT EXISTS providers (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT UNIQUE NOT NULL,
            base_url TEXT NOT NULL,
            api_type TEXT NOT NULL,
            is_active BOOLEAN DEFAULT TRUE
        );

        CREATE TABLE IF NOT EXISTS provider_keys (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            provider_id INTEGER NOT NULL REFERENCES providers(id),
            key TEXT NOT NULL,
            is_active BOOLEAN DEFAULT TRUE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );

        CREATE TABLE IF NOT EXISTS request_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_key_id INTEGER REFERENCES client_keys(id),
            provider_id INTEGER REFERENCES providers(id),
            model TEXT,
            prompt_tokens INTEGER,
            completion_tokens INTEGER,
            latency_ms INTEGER,
            success BOOLEAN,
            error_message TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    """)


async def _seed_test_data(db: aiosqlite.Connection):
    """Seed test data."""
    # Seed providers
    await db.executemany(
        "INSERT INTO providers (name, base_url, api_type) VALUES (?, ?, ?)",
        [
            ("openai", "https://api.openai.com/v1", "openai"),
            ("anthropic", "https://api.anthropic.com/v1", "anthropic"),
        ]
    )
    # Seed client keys
    await db.execute(
        "INSERT INTO client_keys (key, name, is_active) VALUES (?, ?, ?)",
        ("test-api-key", "test-client", True)
    )
    await db.execute(
        "INSERT INTO client_keys (key, name, is_active) VALUES (?, ?, ?)",
        ("disabled-key", "disabled-client", False)
    )
    # Seed provider keys
    await db.execute(
        "INSERT INTO provider_keys (provider_id, key, is_active) VALUES (?, ?, ?)",
        (1, "sk-openai-key-1", True)
    )
    await db.execute(
        "INSERT INTO provider_keys (provider_id, key, is_active) VALUES (?, ?, ?)",
        (1, "sk-openai-key-2", True)
    )
    await db.execute(
        "INSERT INTO provider_keys (provider_id, key, is_active) VALUES (?, ?, ?)",
        (2, "sk-ant-key-1", True)
    )
    await db.commit()


@pytest.fixture
def mock_db(monkeypatch):
    """Mock database connection for tests."""
    @asynccontextmanager
    async def mock_db_connection():
        async with get_test_db() as db:
            yield db

    monkeypatch.setattr("app.database.db_connection", mock_db_connection)
    return mock_db_connection


@pytest.fixture
def mock_litellm():
    """Mock LiteLLM completion."""
    mock_response = MagicMock()
    mock_response.usage = MagicMock()
    mock_response.usage.prompt_tokens = 10
    mock_response.usage.completion_tokens = 20
    mock_response.model_dump.return_value = {
        "id": "chatcmpl-test",
        "object": "chat.completion",
        "choices": [{"message": {"content": "test response"}}]
    }

    async def mock_acompletion(*args, **kwargs):
        return mock_response

    return {"acompletion": mock_acompletion, "response": mock_response}


@pytest.fixture
def client_key():
    """Return a test client key instance."""
    from app.models import ClientKey
    return ClientKey(
        id=1,
        key="test-api-key",
        name="test-client",
        is_active=True,
        created_at="2024-01-01 00:00:00"
    )


@pytest.fixture
def provider():
    """Return a test provider instance."""
    from app.models import Provider
    return Provider(
        id=1,
        name="openai",
        base_url="https://api.openai.com/v1",
        api_type="openai",
        is_active=True
    )


@pytest.fixture
def provider_key():
    """Return a test provider key instance."""
    from app.models import ProviderKey
    return ProviderKey(
        id=1,
        provider_id=1,
        key="sk-test-key",
        is_active=True,
        created_at="2024-01-01 00:00:00"
    )