letsbe-orchestrator/tests/conftest.py

"""Pytest configuration and fixtures for letsbe-orchestrator tests."""

import asyncio
import hashlib
import uuid
from collections.abc import AsyncGenerator

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

from app.models.base import Base, utc_now
from app.models.tenant import Tenant
from app.models.agent import Agent
from app.models.registration_token import RegistrationToken

# Use in-memory SQLite for testing
TEST_DATABASE_URL = "sqlite+aiosqlite:///:memory:"


@pytest.fixture(scope="session")
def event_loop():
    """Create an instance of the default event loop for the test session."""
    loop = asyncio.get_event_loop_policy().new_event_loop()
    yield loop
    loop.close()


@pytest_asyncio.fixture(scope="function")
async def async_engine():
    """Create a test async engine."""
    engine = create_async_engine(TEST_DATABASE_URL, echo=False)
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    yield engine
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)
    await engine.dispose()


@pytest_asyncio.fixture(scope="function")
async def db(async_engine) -> AsyncGenerator[AsyncSession, None]:
    """Create a test database session."""
    session_factory = async_sessionmaker(
        async_engine,
        class_=AsyncSession,
        expire_on_commit=False,
        autocommit=False,
        autoflush=False,
    )
    async with session_factory() as session:
        yield session


@pytest_asyncio.fixture(scope="function")
async def test_tenant(db: AsyncSession) -> Tenant:
    """Create a test tenant."""
    tenant = Tenant(
        id=uuid.uuid4(),
        name="Test Tenant",
    )
    db.add(tenant)
    await db.commit()
    await db.refresh(tenant)
    return tenant


@pytest_asyncio.fixture(scope="function")
async def test_agent(db: AsyncSession) -> Agent:
    """Create a test agent."""
    agent = Agent(
        id=uuid.uuid4(),
        name="test-agent-host",
        version="1.0.0",
        status="online",
        token="test-token-12345",
    )
    db.add(agent)
    await db.commit()
    await db.refresh(agent)
    return agent


@pytest_asyncio.fixture(scope="function")
async def test_agent_for_tenant(db: AsyncSession, test_tenant: Tenant) -> Agent:
    """Create a test agent linked to test_tenant with online status."""
    agent = Agent(
        id=uuid.uuid4(),
        tenant_id=test_tenant.id,
        name="test-agent-for-tenant",
        version="1.0.0",
        status="online",
        token="test-token-tenant",
    )
    db.add(agent)
    await db.commit()
    await db.refresh(agent)
    return agent


# --- New fixtures for secure auth testing ---


@pytest_asyncio.fixture(scope="function")
async def test_registration_token(
    db: AsyncSession, test_tenant: Tenant
) -> tuple[RegistrationToken, str]:
    """Create a test registration token and return (token_record, plaintext_token)."""
    plaintext_token = str(uuid.uuid4())
    token_hash = hashlib.sha256(plaintext_token.encode()).hexdigest()

    reg_token = RegistrationToken(
        id=uuid.uuid4(),
        tenant_id=test_tenant.id,
        token_hash=token_hash,
        description="Test registration token",
        max_uses=10,
        use_count=0,
    )
    db.add(reg_token)
    await db.commit()
    await db.refresh(reg_token)
    return reg_token, plaintext_token


@pytest_asyncio.fixture(scope="function")
async def test_agent_with_secret(
    db: AsyncSession, test_tenant: Tenant
) -> tuple[Agent, str]:
    """Create a test agent with secret_hash and return (agent, plaintext_secret)."""
    plaintext_secret = "test-secret-" + uuid.uuid4().hex[:16]
    secret_hash = hashlib.sha256(plaintext_secret.encode()).hexdigest()

    agent = Agent(
        id=uuid.uuid4(),
        tenant_id=test_tenant.id,
        name="test-agent-secure",
        version="1.0.0",
        status="online",
        token="",  # Empty for new auth scheme
        secret_hash=secret_hash,
        last_heartbeat=utc_now(),
    )
    db.add(agent)
    await db.commit()
    await db.refresh(agent)
    return agent, plaintext_secret


@pytest_asyncio.fixture(scope="function")
async def test_agent_legacy(db: AsyncSession, test_tenant: Tenant) -> Agent:
    """Create a test agent with legacy token auth (both token and secret_hash set)."""
    legacy_token = "legacy-token-" + uuid.uuid4().hex[:16]
    secret_hash = hashlib.sha256(legacy_token.encode()).hexdigest()

    agent = Agent(
        id=uuid.uuid4(),
        tenant_id=test_tenant.id,
        name="test-agent-legacy",
        version="1.0.0",
        status="online",
        token=legacy_token,  # Legacy field
        secret_hash=secret_hash,  # Also set for new auth
    )
    db.add(agent)
    await db.commit()
    await db.refresh(agent)
    return agent