kca/server/tests/test_sqlite.py

"""Persistent SQLite configuration and migration integration tests."""
from __future__ import annotations

import asyncio
import os
import sqlite3
import subprocess
import sys
import tempfile
import unittest
import uuid
from contextlib import closing
from pathlib import Path

from kca.config import _database_url_from_parts
from kca.models import Vault
from kca.tortoise_mixin import TortoiseMixin
from tortoise.context import TortoiseContext

SERVER_ROOT = Path(__file__).resolve().parents[1]
MIGRATIONS_DIR = SERVER_ROOT / "migrations" / "models"


class SqliteConfigurationTests(unittest.TestCase):
    def test_database_url_from_sqlite_parts(self) -> None:
        self.assertEqual(
            _database_url_from_parts(
                "sqlite", "ignored", "ignored", "ignored", "5432", "/data/kca.db", "busy_timeout=5000"
            ),
            "sqlite:///data/kca.db?busy_timeout=5000",
        )
        self.assertEqual(
            _database_url_from_parts("sqlite", "", None, "", "", "data/kca.db", ""),
            "sqlite://data/kca.db",
        )
        self.assertEqual(
            _database_url_from_parts("sqlite", "", None, "", "", ":memory:", ""),
            "sqlite://:memory:",
        )


class SqliteMemoryTests(unittest.TestCase):
    def test_memory_url_with_options_generates_schemas(self) -> None:
        async def create_vault() -> None:
            mixin = TortoiseMixin("sqlite://:memory:?busy_timeout=5000", ["kca.models"])
            await mixin.ensure_context()
            try:
                await Vault.create(
                    id=uuid.uuid4(),
                    owner_issuer="https://issuer.example",
                    owner_sub="memory-user",
                    name="Memory vault",
                    settings={},
                )
                self.assertEqual(await Vault.all().count(), 1)
            finally:
                await mixin.aclose()

        asyncio.run(create_vault())


class SqliteFileMigrationTests(unittest.TestCase):
    def test_migrations_pragmas_and_persistence(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            db_path = Path(tmp) / "kca.db"
            database_url = f"sqlite://{db_path}?busy_timeout=5000"
            env = os.environ.copy()
            env["DATABASE_URL"] = database_url
            aerich = Path(sys.executable).with_name("aerich")

            for _ in range(2):
                result = subprocess.run(
                    [str(aerich), "upgrade"],
                    cwd=SERVER_ROOT,
                    env=env,
                    text=True,
                    capture_output=True,
                    check=False,
                )
                self.assertEqual(result.returncode, 0, result.stdout + result.stderr)

            expected_tables = {
                "aerich",
                "api_token",
                "app_user",
                "authority",
                "certificate",
                "certificate_revocation",
                "crl",
                "key",
                "ocsp_snapshot",
                "request",
                "template",
                "vault",
            }
            migration_count = len(list(MIGRATIONS_DIR.glob("[0-9]_*.py")))
            with closing(sqlite3.connect(db_path)) as db:
                tables = {
                    row[0]
                    for row in db.execute(
                        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
                    )
                }
                self.assertEqual(tables, expected_tables)
                self.assertEqual(db.execute("SELECT COUNT(*) FROM aerich").fetchone()[0], migration_count)
                certificate_columns = {row[1] for row in db.execute("PRAGMA table_info(certificate)")}
                authority_columns = {row[1] for row in db.execute("PRAGMA table_info(authority)")}
                user_columns = {row[1] for row in db.execute("PRAGMA table_info(app_user)")}
                self.assertIn("no_expiration", certificate_columns)
                self.assertIn("template_key", authority_columns)
                self.assertIn("issuer", user_columns)

            vault_id = uuid.uuid4()

            async def write_vault() -> None:
                ctx = TortoiseContext()
                with ctx:
                    await ctx.init(db_url=database_url, modules={"models": ["kca.models"]})
                    try:
                        connection = ctx.connections.get("default")
                        foreign_keys = await connection.execute_query_dict("PRAGMA foreign_keys")
                        journal_mode = await connection.execute_query_dict("PRAGMA journal_mode")
                        busy_timeout = await connection.execute_query_dict("PRAGMA busy_timeout")
                        self.assertEqual(foreign_keys[0]["foreign_keys"], 1)
                        self.assertEqual(journal_mode[0]["journal_mode"].lower(), "wal")
                        self.assertEqual(busy_timeout[0]["timeout"], 5000)
                        await Vault.create(
                            id=vault_id,
                            owner_issuer="https://issuer.example",
                            owner_sub="sqlite-user",
                            name="Persistent vault",
                            settings={},
                        )
                    finally:
                        await ctx.close_connections()

            asyncio.run(write_vault())
            with closing(sqlite3.connect(db_path)) as db:
                row = db.execute("SELECT name FROM vault WHERE id = ?", (str(vault_id),)).fetchone()
                self.assertEqual(row, ("Persistent vault",))

            # SQLite bootstraps at the current schema, so historical downgrades
            # are executable no-ops and preserve both the schema and its data.
            for _ in range(migration_count - 1):
                result = subprocess.run(
                    [str(aerich), "downgrade", "--yes"],
                    cwd=SERVER_ROOT,
                    env=env,
                    text=True,
                    capture_output=True,
                    check=False,
                )
                self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
            result = subprocess.run(
                [str(aerich), "upgrade"],
                cwd=SERVER_ROOT,
                env=env,
                text=True,
                capture_output=True,
                check=False,
            )
            self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
            with closing(sqlite3.connect(db_path)) as db:
                self.assertEqual(db.execute("SELECT COUNT(*) FROM aerich").fetchone()[0], migration_count)
                row = db.execute("SELECT name FROM vault WHERE id = ?", (str(vault_id),)).fetchone()
                self.assertEqual(row, ("Persistent vault",))


if __name__ == "__main__":
    unittest.main()