plm-lottery/tests/unit/test_auth.py

import pytest
from cryptography.fernet import Fernet
from httpx import ASGITransport, AsyncClient

from app.config import settings


@pytest.fixture
async def client(monkeypatch, tmp_path):
    monkeypatch.setattr(settings, "database_url", f"sqlite+aiosqlite:///{tmp_path}/test.db")
    monkeypatch.setattr(settings, "jwt_secret", "test-jwt-secret")
    monkeypatch.setattr(settings, "xprv_encryption_key", Fernet.generate_key().decode())
    monkeypatch.setattr(settings, "master_key_path", str(tmp_path / "master.xprv.enc"))

    import app.wallet.hd as hd

    hd._account_key = None
    hd.generate_master_key()

    from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

    from app.db import base as db_base

    import app.db.models  # noqa: F401

    db_base.engine = create_async_engine(settings.database_url)
    db_base.AsyncSessionLocal = async_sessionmaker(db_base.engine, expire_on_commit=False)

    from app.db import session as db_session

    db_session.AsyncSessionLocal = db_base.AsyncSessionLocal

    async with db_base.engine.begin() as conn:
        await conn.run_sync(db_base.Base.metadata.create_all)

    from fastapi import FastAPI

    from app.auth.routes import router as auth_router
    from app.electrum.listener import ElectrumListener

    app = FastAPI()
    app.include_router(auth_router)
    app.state.electrum_listener = ElectrumListener(lambda endpoint: None, db_base.AsyncSessionLocal)

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

    await db_base.engine.dispose()


async def _register(client, username="alice", password="original-password"):
    resp = await client.post("/auth/register", json={"username": username, "password": password})
    assert resp.status_code == 201
    return resp.json()["access_token"]


async def test_login_locks_out_after_repeated_failures(client):
    await _register(client)

    for _ in range(5):
        resp = await client.post("/auth/login", json={"username": "alice", "password": "wrong"})
        assert resp.status_code == 401

    resp = await client.post("/auth/login", json={"username": "alice", "password": "wrong"})
    assert resp.status_code == 429
    assert resp.json()["detail"]["code"] == "rate_limited"

    # Even the *correct* password is refused while locked out — the throttle
    # protects against a lucky guess landing inside the backoff window too.
    resp = await client.post("/auth/login", json={"username": "alice", "password": "original-password"})
    assert resp.status_code == 429


async def test_unknown_username_and_wrong_password_share_a_bucket_and_response(client):
    await _register(client, username="bob")

    for _ in range(5):
        resp = await client.post("/auth/login", json={"username": "nobody", "password": "wrong"})
        assert resp.status_code == 401
        assert resp.json()["detail"]["code"] == "invalid_credentials"

    resp = await client.post("/auth/login", json={"username": "nobody", "password": "wrong"})
    assert resp.status_code == 429


async def test_login_failures_against_one_account_do_not_lock_out_another(client):
    await _register(client, username="alice")
    await _register(client, username="carol", password="carols-password")

    for _ in range(6):
        await client.post("/auth/login", json={"username": "alice", "password": "wrong"})

    # Different username, but same IP (the test client always looks the same) —
    # only the per-username bucket should be exhausted, not the whole IP, since
    # the per-username threshold (5) is hit well before the shared IP bucket's.
    resp = await client.post("/auth/login", json={"username": "carol", "password": "carols-password"})
    assert resp.status_code == 200


async def test_successful_login_resets_the_username_bucket(client):
    await _register(client)

    for _ in range(4):
        resp = await client.post("/auth/login", json={"username": "alice", "password": "wrong"})
        assert resp.status_code == 401

    resp = await client.post("/auth/login", json={"username": "alice", "password": "original-password"})
    assert resp.status_code == 200

    resp = await client.post("/auth/login", json={"username": "alice", "password": "original-password"})
    assert resp.status_code == 200


async def test_registration_is_rate_limited_per_ip(client):
    for i in range(5):
        resp = await client.post(
            "/auth/register", json={"username": f"user{i}", "password": "a-strong-password"}
        )
        assert resp.status_code == 201

    resp = await client.post(
        "/auth/register", json={"username": "user5", "password": "a-strong-password"}
    )
    assert resp.status_code == 429
    assert resp.json()["detail"]["code"] == "rate_limited"