DealDocumentScreening/tests/integration/test_auth_flow.py

"""Auth integration tests: Telegram identity sources issue a common JWT.

Run against the Docker Compose infrastructure (`docker compose up -d`).
"""

from __future__ import annotations

import hashlib
import hmac
import time
from urllib.parse import urlencode

import httpx
import pytest
from sqlalchemy import text

pytestmark = pytest.mark.integration

_BOT_TOKEN = "it-test-bot-token:it-test-secret"
_SECRET_KEY = hmac.new(
    _BOT_TOKEN.encode("utf-8"),
    b"WebAppData",
    hashlib.sha256,
).digest()


def _make_web_payload(telegram_id: int) -> dict[str, object]:
    auth_date = int(time.time())
    data = {"id": telegram_id, "first_name": "Integration", "auth_date": auth_date}
    data_check_string = "\n".join(f"{k}={v}" for k, v in sorted(data.items()))
    data["hash"] = hmac.new(
        _SECRET_KEY,
        data_check_string.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return data


def _make_init_data(telegram_id: int) -> str:
    auth_date = int(time.time())
    user_json = f'{{"id":{telegram_id},"first_name":"Integration"}}'
    params = {"user": user_json, "auth_date": str(auth_date)}
    data_check_string = "\n".join(f"{k}={v}" for k, v in sorted(params.items()))
    params["hash"] = hmac.new(
        _SECRET_KEY,
        data_check_string.encode("utf-8"),
        hashlib.sha256,
    ).hexdigest()
    return urlencode(params)


async def test_auth_telegram_bot_creates_user_and_issues_jwt(
    client: httpx.AsyncClient,
    infra: dict[str, str],
    db_session,
) -> None:
    telegram_id = 123_456_789

    r = await client.post(
        "/api/v1/auth/telegram/bot",
        json={"telegram_id": telegram_id},
        headers={"Authorization": f"Bearer {infra['token']}"},
    )
    assert r.status_code == 200
    body = r.json()
    assert body["token_type"] == "bearer"
    assert body["telegram_id"] == telegram_id
    assert body["access_token"]

    # User row created.
    result = await db_session.execute(
        text("SELECT id, telegram_id FROM users WHERE telegram_id = :t"),
        {"t": telegram_id},
    )
    row = result.first()
    assert row is not None
    assert str(row[0]) == body["user_id"]

    # JWT works on /me.
    me = await client.get(
        "/api/v1/me",
        headers={"Authorization": f"Bearer {body['access_token']}"},
    )
    assert me.status_code == 200
    assert me.json()["telegram_id"] == telegram_id


async def test_auth_telegram_web_verifies_widget_payload(
    client: httpx.AsyncClient,
    infra: dict[str, str],
) -> None:
    telegram_id = 222_333_444

    r = await client.post(
        "/api/v1/auth/telegram/web",
        json=_make_web_payload(telegram_id),
    )
    assert r.status_code == 200
    body = r.json()
    assert body["telegram_id"] == telegram_id
    assert body["access_token"]


async def test_auth_telegram_miniapp_verifies_init_data(
    client: httpx.AsyncClient,
    infra: dict[str, str],
) -> None:
    telegram_id = 333_444_555

    r = await client.post(
        "/api/v1/auth/telegram/miniapp",
        json={"init_data": _make_init_data(telegram_id)},
    )
    assert r.status_code == 200
    body = r.json()
    assert body["telegram_id"] == telegram_id
    assert body["access_token"]


async def test_auth_telegram_web_rejects_bad_hash(
    client: httpx.AsyncClient,
    infra: dict[str, str],
) -> None:
    payload = _make_web_payload(444_555_666)
    payload["hash"] = "0" * 64

    r = await client.post("/api/v1/auth/telegram/web", json=payload)
    assert r.status_code == 401


async def test_protected_endpoint_rejects_missing_token(
    client: httpx.AsyncClient,
) -> None:
    r = await client.get("/api/v1/me")
    assert r.status_code == 401


async def test_protected_endpoint_rejects_service_token(
    client: httpx.AsyncClient,
    infra: dict[str, str],
) -> None:
    r = await client.get(
        "/api/v1/me",
        headers={"Authorization": f"Bearer {infra['token']}"},
    )
    assert r.status_code == 401


async def test_auth_me_introspects_jwt(
    client: httpx.AsyncClient,
    infra: dict[str, str],
) -> None:
    telegram_id = 555_666_777
    r = await client.post(
        "/api/v1/auth/telegram/bot",
        json={"telegram_id": telegram_id},
        headers={"Authorization": f"Bearer {infra['token']}"},
    )
    token = r.json()["access_token"]

    introspect = await client.get(
        "/api/v1/auth/me",
        headers={"Authorization": f"Bearer {token}"},
    )
    assert introspect.status_code == 200
    body = introspect.json()
    assert body["telegram_id"] == telegram_id
    assert body["type"] == "access"