YPMS/tests/auth/test_router.py

"""End-to-end auth flow via the FastAPI TestClient.

Uses the shared SQLite in-memory DB by wiring the `get_db` dependency to the
per-test session, then seeding a user and driving login/refresh/me/
change-password.
"""

from __future__ import annotations

from collections.abc import Iterator

import pytest
from fastapi.testclient import TestClient
from sqlalchemy.orm import Session

from ypms.auth.rate_limit import login_limiter
from ypms.auth.security import (
    create_access_token,
    create_refresh_token,
    hash_password,
)
from ypms.db import get_db
from ypms.main import create_app
from ypms.security.rbac import ADMIN_ROLE
from ypms.users.models import (
    ASSURANCE_ADMIN,
    ASSURANCE_STANDARD,
    Role,
    User,
    UserRole,
)


@pytest.fixture(autouse=True)
def _reset_rate_limits() -> Iterator[None]:
    login_limiter.reset()
    yield
    login_limiter.reset()


@pytest.fixture
def staff_user(db: Session) -> User:
    user = User(
        username="staff.user",
        email="staff@dc.gov",
        full_name="Staff User",
        password_hash=hash_password("correct-horse-battery"),
        is_active=True,
        assurance_level=ASSURANCE_STANDARD,
    )
    db.add(user)
    db.flush()
    db.refresh(user)
    return user


@pytest.fixture
def admin_user(db: Session) -> User:
    role = Role(code=ADMIN_ROLE, description="super-role")
    db.add(role)
    db.flush()
    user = User(
        username="chris",
        email="chris@dc.gov",
        full_name="Chris Tonjes",
        password_hash=hash_password("correct-horse-battery"),
        is_active=True,
        assurance_level=ASSURANCE_STANDARD,  # role bump upgrades to admin
    )
    db.add(user)
    db.flush()
    db.add(UserRole(user_id=user.id, role_id=role.id))
    db.flush()
    db.refresh(user)
    return user


@pytest.fixture
def app_client(db: Session) -> Iterator[TestClient]:
    """A TestClient wired to the per-test SQLite session via dep override."""
    app = create_app()

    def _override_get_db() -> Iterator[Session]:
        yield db

    app.dependency_overrides[get_db] = _override_get_db
    with TestClient(app) as c:
        yield c
    app.dependency_overrides.clear()


def test_login_wrong_password_returns_401(app_client: TestClient, staff_user: User) -> None:
    r = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "wrong"},
    )
    assert r.status_code == 401
    assert r.json()["detail"] == "invalid credentials"


def test_login_success_returns_tokens(app_client: TestClient, staff_user: User) -> None:
    r = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    assert r.status_code == 200
    body = r.json()
    assert body["token_type"] == "bearer"
    assert body["access_token"]
    assert body["refresh_token"]


def test_admin_login_upgrades_assurance_level(
    app_client: TestClient, admin_user: User
) -> None:
    from ypms.auth.security import decode_token

    r = app_client.post(
        "/auth/login",
        data={"username": admin_user.username, "password": "correct-horse-battery"},
    )
    assert r.status_code == 200
    claims = decode_token(r.json()["access_token"])
    assert claims["iam.assurance_level"] == ASSURANCE_ADMIN
    assert ADMIN_ROLE in claims["roles"]


def test_me_requires_authentication(app_client: TestClient) -> None:
    r = app_client.get("/auth/me")
    assert r.status_code == 401


def test_me_returns_current_user(app_client: TestClient, staff_user: User) -> None:
    login = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    access = login.json()["access_token"]
    r = app_client.get("/auth/me", headers={"Authorization": f"Bearer {access}"})
    assert r.status_code == 200
    assert r.json()["username"] == staff_user.username
    assert r.json()["email"] == "staff@dc.gov"


def test_refresh_returns_new_access_token(app_client: TestClient, staff_user: User) -> None:
    login = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    refresh = login.json()["refresh_token"]
    r = app_client.post("/auth/refresh", json={"refresh_token": refresh})
    assert r.status_code == 200
    assert r.json()["access_token"]
    assert "refresh_token" not in r.json()


def test_refresh_rejects_access_token(app_client: TestClient, staff_user: User) -> None:
    access = create_access_token(
        sub=staff_user.username,
        uid=staff_user.id,
        roles=[],
        assurance_level=ASSURANCE_STANDARD,
        mfa_verified=False,
    )
    r = app_client.post("/auth/refresh", json={"refresh_token": access})
    assert r.status_code == 401


def test_me_rejects_refresh_token(app_client: TestClient, staff_user: User) -> None:
    refresh = create_refresh_token(sub=staff_user.username, uid=staff_user.id)
    r = app_client.get("/auth/me", headers={"Authorization": f"Bearer {refresh}"})
    assert r.status_code == 401


def test_change_password_flow(app_client: TestClient, staff_user: User) -> None:
    login = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    access = login.json()["access_token"]
    r = app_client.post(
        "/auth/change-password",
        headers={"Authorization": f"Bearer {access}"},
        json={
            "current_password": "correct-horse-battery",
            "new_password": "totally-different-passphrase",
        },
    )
    assert r.status_code == 204

    # Old password no longer works.
    r2 = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    assert r2.status_code == 401

    # New password does.
    r3 = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "totally-different-passphrase"},
    )
    assert r3.status_code == 200


def test_change_password_rejects_short_new_password(
    app_client: TestClient, staff_user: User
) -> None:
    login = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    access = login.json()["access_token"]
    r = app_client.post(
        "/auth/change-password",
        headers={"Authorization": f"Bearer {access}"},
        json={"current_password": "correct-horse-battery", "new_password": "short"},
    )
    assert r.status_code == 400


def test_change_password_rejects_wrong_current(
    app_client: TestClient, staff_user: User
) -> None:
    login = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    access = login.json()["access_token"]
    r = app_client.post(
        "/auth/change-password",
        headers={"Authorization": f"Bearer {access}"},
        json={"current_password": "wrong", "new_password": "long-enough-password"},
    )
    assert r.status_code == 401


def test_login_lockout_after_five_failed_attempts(
    app_client: TestClient, staff_user: User
) -> None:
    for _ in range(5):
        r = app_client.post(
            "/auth/login",
            data={"username": staff_user.username, "password": "wrong"},
        )
        assert r.status_code == 401
    # 6th attempt (any credentials) is rate-limited.
    r = app_client.post(
        "/auth/login",
        data={"username": staff_user.username, "password": "correct-horse-battery"},
    )
    assert r.status_code == 429
    assert r.headers["Retry-After"] == "900"