ctms/backend/tests/test_registration.py

import pytest
import pytest_asyncio
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding
from httpx import AsyncClient
from sqlalchemy import UUID as SA_UUID, select
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
import base64
import json
import os

from app.main import create_app
from app.core.deps import get_db_session
from app.core.security import hash_password
from app.crud import user as user_crud
from app.db.base_class import Base
from tests.conftest import GUID
from app.models.study_member import StudyMember
from app.models.user import User, UserRole, UserStatus
from app.schemas.user import UserRegisterRequest

TEST_DATABASE_URL = "sqlite+aiosqlite:///:memory:"


async def encrypted_auth_payload(client: AsyncClient, email: str, password: str) -> dict:
    key_resp = await client.get("/api/v1/auth/login-key")
    assert key_resp.status_code == 200
    login_key = key_resp.json()
    public_key = serialization.load_pem_public_key(login_key["public_key"].encode("utf-8"))
    plaintext = json.dumps(
        {
            "email": email,
            "password": password,
            "challenge": login_key["challenge"],
        },
        separators=(",", ":"),
    ).encode("utf-8")
    aes_key = AESGCM.generate_key(bit_length=256)
    iv = os.urandom(12)
    encrypted_data = AESGCM(aes_key).encrypt(iv, plaintext, None)
    encrypted_key = public_key.encrypt(
        aes_key,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return {
        "key_id": login_key["key_id"],
        "challenge": login_key["challenge"],
        "ciphertext": base64.b64encode(
            json.dumps(
                {
                    "encrypted_key": base64.b64encode(encrypted_key).decode("ascii"),
                    "iv": base64.b64encode(iv).decode("ascii"),
                    "data": base64.b64encode(encrypted_data).decode("ascii"),
                },
                separators=(",", ":"),
            ).encode("utf-8")
        ).decode("ascii"),
    }


async def encrypted_login(client: AsyncClient, email: str, password: str):
    return await client.post("/api/v1/auth/login", json=await encrypted_auth_payload(client, email, password))


@pytest_asyncio.fixture
async def client_and_db():
    engine = create_async_engine(TEST_DATABASE_URL, future=True)
    TestSessionLocal = async_sessionmaker(engine, expire_on_commit=False)
    original_column_types = {
        column: column.type
        for table in Base.metadata.tables.values()
        for column in table.columns
    }
    for table in Base.metadata.tables.values():
        for column in table.columns:
            if isinstance(column.type, (PG_UUID, SA_UUID)):
                column.type = GUID()
    try:
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

        async def override_get_db():
            async with TestSessionLocal() as session:
                yield session

        app = create_app()
        app.dependency_overrides[get_db_session] = override_get_db

        async with TestSessionLocal() as session:
            admin = User(
                email="admin@test.com",
                password_hash=hash_password("admin123"),
                full_name="Admin",
                clinical_department="Admin",
                role=UserRole.ADMIN,
                status=UserStatus.ACTIVE,
            )
            session.add(admin)
            await session.commit()

        async with AsyncClient(app=app, base_url="http://test") as client:
            yield client, TestSessionLocal
    finally:
        for column, original_type in original_column_types.items():
            column.type = original_type
        await engine.dispose()


@pytest.mark.asyncio
async def test_register_creates_pending_user(client_and_db):
    client, SessionLocal = client_and_db
    payload = {
        "email": "newuser@test.com",
        "password": "Password123",
        "full_name": "New User",
        "clinical_department": "Clinical",
    }
    resp = await client.post("/api/v1/auth/register", json=payload)
    assert resp.status_code == 201
    async with SessionLocal() as session:
        user = await user_crud.get_by_email(session, payload["email"])
        assert user is not None
        assert user.status == UserStatus.PENDING
        assert user.role == UserRole.PV
        member_rows = (
            await session.execute(select(StudyMember).where(StudyMember.user_id == user.id))
        ).scalars().all()
        assert member_rows == []


@pytest.mark.asyncio
async def test_root_returns_service_metadata(client_and_db):
    client, _ = client_and_db
    resp = await client.get("/")
    assert resp.status_code == 200
    assert resp.json() == {
        "service": "ctms-backend",
        "status": "ok",
        "health": "/health",
        "docs": "/docs",
        "api": "/api/v1",
    }


@pytest.mark.asyncio
async def test_login_blocked_before_approval(client_and_db):
    client, SessionLocal = client_and_db
    payload = {
        "email": "pending@test.com",
        "password": "Password123",
        "full_name": "Pending User",
        "clinical_department": "Safety",
    }
    await client.post("/api/v1/auth/register", json=payload)
    resp = await encrypted_login(client, payload["email"], payload["password"])
    assert resp.status_code == 401
    assert "账号未审核" in resp.json().get("detail", "")


@pytest.mark.asyncio
async def test_admin_can_approve_user(client_and_db):
    client, SessionLocal = client_and_db
    payload = {
        "email": "approve@test.com",
        "password": "Password123",
        "full_name": "Approve Target",
        "clinical_department": "Supply",
    }
    await client.post("/api/v1/auth/register", json=payload)
    async with SessionLocal() as session:
        user = await user_crud.get_by_email(session, payload["email"])
        user_id = user.id

    admin_login = await encrypted_login(client, "admin@test.com", "admin123")
    token = admin_login.json()["access_token"]
    headers = {"Authorization": f"Bearer {token}"}

    resp = await client.post(f"/api/v1/admin/users/{user_id}/approve", json={"action": "approve"}, headers=headers)
    assert resp.status_code == 200
    async with SessionLocal() as session:
        refreshed = await user_crud.get_by_email(session, payload["email"])
        assert refreshed.status == UserStatus.ACTIVE
        assert refreshed.approved_by is not None


def test_register_request_does_not_expose_role_input():
    assert "role" not in UserRegisterRequest.model_fields


@pytest.mark.asyncio
async def test_plaintext_login_is_rejected(client_and_db):
    client, _ = client_and_db
    resp = await client.post("/api/v1/auth/login", json={"email": "admin@test.com", "password": "admin123"})
    assert resp.status_code == 422


@pytest.mark.asyncio
async def test_login_challenge_cannot_be_reused(client_and_db):
    client, _ = client_and_db
    key_resp = await client.get("/api/v1/auth/login-key")
    login_key = key_resp.json()
    public_key = serialization.load_pem_public_key(login_key["public_key"].encode("utf-8"))
    plaintext = json.dumps(
        {
            "email": "admin@test.com",
            "password": "admin123",
            "challenge": login_key["challenge"],
        },
        separators=(",", ":"),
    ).encode("utf-8")
    aes_key = AESGCM.generate_key(bit_length=256)
    iv = os.urandom(12)
    encrypted_data = AESGCM(aes_key).encrypt(iv, plaintext, None)
    encrypted_key = public_key.encrypt(
        aes_key,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        )
    )
    ciphertext = base64.b64encode(
        json.dumps(
            {
                "encrypted_key": base64.b64encode(encrypted_key).decode("ascii"),
                "iv": base64.b64encode(iv).decode("ascii"),
                "data": base64.b64encode(encrypted_data).decode("ascii"),
            },
            separators=(",", ":"),
        ).encode("utf-8")
    ).decode("ascii")
    payload = {
        "key_id": login_key["key_id"],
        "challenge": login_key["challenge"],
        "ciphertext": ciphertext,
    }

    first = await client.post("/api/v1/auth/login", json=payload)
    second = await client.post("/api/v1/auth/login", json=payload)

    assert first.status_code == 200
    assert second.status_code == 401


@pytest.mark.asyncio
async def test_unlock_requires_encrypted_password(client_and_db):
    client, _ = client_and_db
    plaintext = await client.post("/api/v1/auth/unlock", json={"email": "admin@test.com", "password": "admin123"})
    encrypted = await client.post(
        "/api/v1/auth/unlock",
        json=await encrypted_auth_payload(client, "admin@test.com", "admin123"),
    )

    assert plaintext.status_code == 422
    assert encrypted.status_code == 200