opsxsql/backend/tests/conftest.py

from typing import AsyncGenerator

import pytest_asyncio
from httpx import ASGITransport, AsyncClient
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy import text

from app.core.database import Base, get_db
from app.core.security import hash_password, create_access_token
from app.models.user import User
from app.models.instance import Instance
from app.models.instance_group import InstanceGroup, user_groups
from app.models.ticket import Ticket
from app.models.audit_log import AuditLog
from app.models.audit_rule import AuditRule
from app.models.announcement import Announcement
from app.models.notification import NotificationChannel
from app.models.ldap_config import LdapConfig
from app.models.backup import Backup

SQLALCHEMY_DATABASE_URL = "sqlite+aiosqlite://"

engine = create_async_engine(SQLALCHEMY_DATABASE_URL, echo=False)
TestSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


@pytest_asyncio.fixture(scope="session", autouse=True)
async def setup_database():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)


@pytest_asyncio.fixture(autouse=True)
async def clean_tables():
    yield
    async with engine.begin() as conn:
        for table in reversed(Base.metadata.sorted_tables):
            await conn.execute(text(f"DELETE FROM {table.name}"))


@pytest_asyncio.fixture
async def db() -> AsyncGenerator[AsyncSession, None]:
    async with TestSessionLocal() as session:
        yield session


@pytest_asyncio.fixture
async def client(db: AsyncSession) -> AsyncGenerator[AsyncClient, None]:
    from app.main import app

    async def _override():
        yield db

    app.dependency_overrides[get_db] = _override

    transport = ASGITransport(app=app)
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c

    app.dependency_overrides.clear()


@pytest_asyncio.fixture
async def admin_user(db: AsyncSession) -> User:
    user = User(
        username="admin",
        display_name="管理员",
        role="admin",
        hashed_password=hash_password("123456"),
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


@pytest_asyncio.fixture
async def dba_user(db: AsyncSession) -> User:
    user = User(
        username="dba",
        display_name="DBA张三",
        role="dba",
        hashed_password=hash_password("123456"),
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


@pytest_asyncio.fixture
async def developer_user(db: AsyncSession) -> User:
    user = User(
        username="dev1",
        display_name="开发者李四",
        role="developer",
        hashed_password=hash_password("123456"),
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


@pytest_asyncio.fixture
async def guest_user(db: AsyncSession) -> User:
    user = User(
        username="guest1",
        display_name="访客",
        role="guest",
        hashed_password=hash_password("123456"),
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


@pytest_asyncio.fixture
async def disabled_user(db: AsyncSession) -> User:
    user = User(
        username="disabled",
        display_name="已禁用用户",
        role="developer",
        is_active=False,
        hashed_password=hash_password("123456"),
    )
    db.add(user)
    await db.commit()
    await db.refresh(user)
    return user


@pytest_asyncio.fixture
async def admin_token(admin_user: User) -> str:
    return create_access_token(data={"sub": str(admin_user.id), "role": admin_user.role})


@pytest_asyncio.fixture
async def dba_token(dba_user: User) -> str:
    return create_access_token(data={"sub": str(dba_user.id), "role": dba_user.role})


@pytest_asyncio.fixture
async def developer_token(developer_user: User) -> str:
    return create_access_token(data={"sub": str(developer_user.id), "role": developer_user.role})


@pytest_asyncio.fixture
async def guest_token(guest_user: User) -> str:
    return create_access_token(data={"sub": str(guest_user.id), "role": guest_user.role})


def auth_header(token: str) -> dict:
    return {"Authorization": f"Bearer {token}"}


@pytest_asyncio.fixture
async def test_instance(db: AsyncSession) -> Instance:
    instance = Instance(
        name="测试MySQL",
        db_type="mysql",
        host="127.0.0.1",
        port=3306,
        username="root",
        password="test123",
        database="test_db",
        env="development",
        status="disconnected",
    )
    db.add(instance)
    await db.commit()
    await db.refresh(instance)
    return instance