headquarter/apps/api/tests/integration/test_models.py

from collections.abc import AsyncIterator

import pytest
import pytest_asyncio
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from src.config import build_database_url
from src.models import Base
from src.models.base import TimestampMixin, UUIDPrimaryKeyMixin
from src.models.git_repository import GitRepository
from src.models.project import Project
from src.models.refresh_token import RefreshToken
from src.models.ssh_key import SSHKey
from src.models.user import User
from src.models.user_config import UserConfig


TEST_DATABASE_URL = build_database_url(
    user="headquarter",
    password="headquarter",
    host="localhost",
    port=5432,
    database="headquarter",
)


@pytest_asyncio.fixture
async def db_session() -> AsyncIterator[AsyncSession]:
    engine = create_async_engine(TEST_DATABASE_URL)
    session_factory = async_sessionmaker(engine, expire_on_commit=False)

    async with session_factory() as session:
        table_rows = await session.execute(
            text(
                "SELECT tablename FROM pg_tables "
                "WHERE schemaname = 'public' "
                "AND tablename = ANY(:table_names)"
            ),
            {
                "table_names": [
                    "refresh_tokens",
                    "user_configs",
                    "git_repositories",
                    "projects",
                    "ssh_keys",
                    "users",
                ]
            },
        )
        existing_tables = [row[0] for row in table_rows]
        if existing_tables:
            await session.execute(text(f"TRUNCATE TABLE {', '.join(existing_tables)} RESTART IDENTITY CASCADE"))
        await session.commit()
        yield session
        await session.rollback()

    await engine.dispose()


n@pytest.mark.integration

def test_base_metadata_collects_declared_tables() -> None:
    assert isinstance(Base.metadata.tables, dict)


n@pytest.mark.integration

def test_shared_mixins_define_expected_columns() -> None:
    assert "id" in UUIDPrimaryKeyMixin.__dict__
    assert "created_at" in TimestampMixin.__dict__
    assert "updated_at" in TimestampMixin.__dict__


n@pytest.mark.integration

def test_expected_tables_are_registered() -> None:
    assert set(Base.metadata.tables) == {
        "refresh_tokens",
        "git_repositories",
        "projects",
        "ssh_keys",
        "user_configs",
        "users",
    }


n@pytest.mark.integration

def test_user_table_has_required_columns() -> None:
    columns = User.__table__.columns

    assert set(columns.keys()) == {
        "id",
        "email",
        "name",
        "authentik_id",
        "avatar_url",
        "created_at",
        "updated_at",
    }
    assert columns["email"].unique is True
    assert columns["authentik_id"].unique is True
    assert columns["avatar_url"].nullable is True


n@pytest.mark.integration

def test_project_relationships_point_to_owner_and_default_ssh_key() -> None:
    owner_fk = next(iter(Project.__table__.c.owner_id.foreign_keys))
    ssh_fk = next(iter(Project.__table__.c.default_ssh_key_id.foreign_keys))

    assert owner_fk.target_fullname == "users.id"
    assert ssh_fk.target_fullname == "ssh_keys.id"
    assert Project.owner.property.mapper.class_ is User
    assert Project.default_ssh_key.property.mapper.class_ is SSHKey


n@pytest.mark.integration

def test_repository_and_user_config_relationships_are_registered() -> None:
    project_fk = next(iter(GitRepository.__table__.c.project_id.foreign_keys))
    owner_fk = next(iter(GitRepository.__table__.c.owner_id.foreign_keys))
    user_config_fk = next(iter(UserConfig.__table__.c.user_id.foreign_keys))

    assert project_fk.target_fullname == "projects.id"
    assert owner_fk.target_fullname == "users.id"
    assert user_config_fk.target_fullname == "users.id"
    assert GitRepository.project.property.mapper.class_ is Project
    assert GitRepository.owner.property.mapper.class_ is User
    assert UserConfig.user.property.mapper.class_ is User


n@pytest.mark.integration

def test_refresh_token_table_has_required_columns_and_relationships() -> None:
    columns = RefreshToken.__table__.columns
    user_fk = next(iter(RefreshToken.__table__.c.user_id.foreign_keys))

    assert set(columns.keys()) == {
        "id",
        "user_id",
        "token_hash",
        "expires_at",
        "revoked_at",
        "user_agent",
        "ip_address",
        "created_at",
    }
    assert columns["token_hash"].unique is True
    assert columns["revoked_at"].nullable is True
    assert user_fk.target_fullname == "users.id"
    assert RefreshToken.user.property.mapper.class_ is User


@pytest.mark.asyncio
n@pytest.mark.integration

async def test_async_session_can_insert_and_load_user(db_session: AsyncSession) -> None:
    user = User(email="dev@headquarter.local", name="Dev User", authentik_id="dev-user", avatar_url=None)

    db_session.add(user)
    await db_session.commit()
    await db_session.refresh(user)

    loaded_user = await db_session.get(User, user.id)

    assert loaded_user is not None
    assert loaded_user.email == "dev@headquarter.local"