zzrouter/backend/tests/test_database.py

"""Tests for database module."""
import pytest
import aiosqlite
from unittest.mock import patch, AsyncMock
from app.database import get_db, db_connection, init_db, seed_default_providers


class TestDatabase:
    """Test database functions."""

    @pytest.mark.asyncio
    async def test_get_db_returns_connection(self):
        """Test get_db returns aiosqlite connection."""
        conn = await get_db()
        assert isinstance(conn, aiosqlite.Connection)
        await conn.close()

    @pytest.mark.asyncio
    async def test_get_db_row_factory(self):
        """Test get_db sets row factory."""
        conn = await get_db()
        assert conn.row_factory == aiosqlite.Row
        await conn.close()

    @pytest.mark.asyncio
    async def test_db_connection_context_manager(self):
        """Test db_connection context manager."""
        async with db_connection() as conn:
            assert isinstance(conn, aiosqlite.Connection)
            # Connection should be open
            cursor = await conn.execute("SELECT 1")
            result = await cursor.fetchone()
            assert result is not None

    @pytest.mark.asyncio
    async def test_db_connection_closes_after_context(self):
        """Test connection is closed after context exits."""
        async with db_connection() as conn:
            pass
        # Connection should be closed now
        # This will raise if connection is still open
        with pytest.raises(Exception):
            await conn.execute("SELECT 1")

    @pytest.mark.asyncio
    async def test_init_db_creates_tables(self):
        """Test init_db creates all required tables."""
        # Use in-memory database for testing
        async with aiosqlite.connect(":memory:") as db:
            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
                );
            """)
            await db.commit()

            # Verify tables exist
            cursor = await db.execute(
                "SELECT name FROM sqlite_master WHERE type='table'"
            )
            tables = [row[0] async for row in cursor]
            assert "client_keys" in tables
            assert "providers" in tables
            assert "provider_keys" in tables
            assert "request_logs" in tables

    @pytest.mark.asyncio
    async def test_seed_default_providers(self):
        """Test seed_default_providers inserts OpenAI and Anthropic."""
        async with aiosqlite.connect(":memory:") as db:
            db.row_factory = aiosqlite.Row
            # Create providers table
            await db.execute("""
                CREATE TABLE 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
                )
            """)
            await db.commit()

            # 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"),
                ]
            )
            await db.commit()

            # Verify providers
            cursor = await db.execute("SELECT name, api_type FROM providers")
            providers = {row["name"]: row["api_type"] async for row in cursor}
            assert providers["openai"] == "openai"
            assert providers["anthropic"] == "anthropic"

    @pytest.mark.asyncio
    async def test_seed_default_providers_idempotent(self):
        """Test seeding is idempotent (INSERT OR IGNORE)."""
        async with aiosqlite.connect(":memory:") as db:
            db.row_factory = aiosqlite.Row
            await db.execute("""
                CREATE TABLE 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
                )
            """)
            await db.commit()

            # Insert twice
            for _ in range(2):
                await db.executemany(
                    "INSERT OR IGNORE INTO providers (name, base_url, api_type) VALUES (?, ?, ?)",
                    [
                        ("openai", "https://api.openai.com/v1", "openai"),
                        ("anthropic", "https://api.anthropic.com/v1", "anthropic"),
                    ]
                )
                await db.commit()

            cursor = await db.execute("SELECT COUNT(*) FROM providers")
            count = (await cursor.fetchone())[0]
            assert count == 2  # Should still be 2, not 4