model-platform/backend/tests/test_jupyter_auth_cache.py

"""Unit tests for the 5s auth-result cache in ``backend.api.jupyter``.

Jupyter 一次会话会触发几十次 Nginx ``auth_request``；本缓存按
``(workspace_id, user_id)`` 缓存 membership + runtime 的查找结果，避免
每次都跑 DB JOIN 与跨进程 RPC。JWT 验签与 per-URI 的 lock check **不进
缓存**，每请求都执行。

这些测试直接调用 ``verify_jupyter_access``（不经过 FastAPI TestClient），
用 SimpleNamespace 构造 fake request / response / runtime，并用
monkeypatch 替换 JWT / membership / lock / runtime 的调用点来统计次数。
"""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock

import pytest
from fastapi import HTTPException

import backend.api.jupyter as jupyter_module
from backend.api.jupyter import (
    _JUPYTER_AUTH_CACHE,
    verify_jupyter_access,
)
from backend.clients.runtime import RuntimeClientError

WS_ID = "01WS0000000000000000000A"
USER_ID = "01USR0000000000000000000A"
NOTEBOOK_URI = f"/jupyter/{WS_ID}/notebooks/a.ipynb"
NON_NOTEBOOK_URI = f"/jupyter/{WS_ID}/tree"

_DESCRIPTOR = {
    "status": "running",
    "workspace_id": WS_ID,
    "base_url": "http://runtime",
    "port": 34567,
    "token": "jupyter-token",
}


@pytest.fixture(autouse=True)
def _clear_cache() -> None:
    _JUPYTER_AUTH_CACHE.clear()
    yield
    _JUPYTER_AUTH_CACHE.clear()


def _make_context(runtime_client, uri: str = NOTEBOOK_URI) -> tuple[SimpleNamespace, SimpleNamespace]:
    request = SimpleNamespace(
        headers={
            "X-Original-Workspace-Id": WS_ID,
            "X-Original-URI": uri,
        },
        cookies={"access_token": "a.b.c"},
        app=SimpleNamespace(state=SimpleNamespace(runtime_client=runtime_client)),
    )
    response = SimpleNamespace(headers={})
    return request, response


def _make_runtime_client(descriptor=None) -> SimpleNamespace:
    client = SimpleNamespace()
    client.get_workspace = AsyncMock(return_value=descriptor if descriptor is not None else _DESCRIPTOR)
    client.start_workspace = AsyncMock(return_value=_DESCRIPTOR)
    return client


def _setup_mocks(
    monkeypatch: pytest.MonkeyPatch,
    *,
    user_id: str = USER_ID,
    locked: bool = False,
    runtime_client: SimpleNamespace | None = None,
) -> tuple[SimpleNamespace, SimpleNamespace, AsyncMock, AsyncMock, AsyncMock, SimpleNamespace]:
    """Patch the call points once and return fakes for counting.

    ``verify_jwt_token`` 默认替换为固定 payload；需要计数的测试可在此之后
    再次 ``monkeypatch.setattr`` 覆盖（后设置者生效）。
    """
    monkeypatch.setattr(
        jupyter_module,
        "verify_jwt_token",
        lambda _token: {"sub": user_id},
    )
    membership = AsyncMock()
    monkeypatch.setattr(jupyter_module, "load_active_membership_or_403", membership)
    lock_check = AsyncMock(return_value=locked)
    monkeypatch.setattr(jupyter_module, "check_notebook_is_locked", lock_check)
    runtime = runtime_client if runtime_client is not None else _make_runtime_client()
    request, response = _make_context(runtime)
    return request, response, membership, lock_check, runtime


async def _call_once(request, response) -> None:
    await verify_jupyter_access(
        request,
        response,
        auth=None,
        session=AsyncMock(),
    )


async def test_cache_hit_skips_membership_and_runtime(monkeypatch) -> None:
    """第一次跑完整路径，第二次同样 (ws, user) 跳过 membership + runtime。"""
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)

    await _call_once(request, response)
    assert membership.await_count == 1
    assert runtime.get_workspace.await_count == 1
    assert runtime.start_workspace.await_count == 0
    assert response.headers["x-upstream-addr"] == "http://runtime:34567"
    assert response.headers["x-jupyter-internal-token"] == "jupyter-token"

    # 第二次请求：缓存命中，membership / runtime 不再执行。
    response.headers = {}
    await _call_once(request, response)
    assert membership.await_count == 1
    assert runtime.get_workspace.await_count == 1
    assert runtime.start_workspace.await_count == 0
    # lock check 每请求都跑。
    assert lock_check.await_count == 2
    # 缓存命中也要写 headers。
    assert response.headers["x-upstream-addr"] == "http://runtime:34567"
    assert response.headers["x-jupyter-internal-token"] == "jupyter-token"


async def test_cache_miss_runs_full_path(monkeypatch) -> None:
    """清空缓存后第一次请求必须跑 membership + runtime。"""
    _JUPYTER_AUTH_CACHE.clear()
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)

    await _call_once(request, response)

    assert membership.await_count == 1
    assert runtime.get_workspace.await_count == 1
    assert lock_check.await_count == 1


async def test_lock_check_runs_every_request_even_on_cache_hit(monkeypatch) -> None:
    """缓存命中时仍要执行 lock check（per-URI，不进缓存）。"""
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)

    await _call_once(request, response)  # 预热缓存
    response.headers = {}
    await _call_once(request, response)  # 缓存命中

    assert membership.await_count == 1
    assert lock_check.await_count == 2


async def test_jwt_verify_runs_every_request(monkeypatch) -> None:
    """JWT 验签是每请求的安全边界，缓存命中也不能跳过。"""
    verify_calls: list[int] = []

    def _fake_verify(_token):
        verify_calls.append(1)
        return {"sub": USER_ID}

    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)
    monkeypatch.setattr(jupyter_module, "verify_jwt_token", _fake_verify)

    await _call_once(request, response)
    response.headers = {}
    await _call_once(request, response)  # 缓存命中

    assert membership.await_count == 1
    assert len(verify_calls) == 2


async def test_cache_ttl_expires_after_5s(monkeypatch) -> None:
    """TTL 用 time.monotonic；5s 后缓存失效，重新走完整路径。"""
    now = [100.0]
    monkeypatch.setattr("time.monotonic", lambda: now[0])
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)

    await _call_once(request, response)  # t=100，写缓存（expires=105）
    assert membership.await_count == 1

    response.headers = {}
    await _call_once(request, response)  # t=100，缓存命中
    assert membership.await_count == 1

    now[0] = 105.0  # 恰好到过期时刻 -> 缓存失效
    response.headers = {}
    await _call_once(request, response)
    assert membership.await_count == 2
    assert runtime.get_workspace.await_count == 2


async def test_lock_check_failure_does_not_populate_cache(monkeypatch) -> None:
    """lock 403 不进缓存。"""
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch, locked=True)

    with pytest.raises(HTTPException) as excinfo:
        await _call_once(request, response)

    assert excinfo.value.status_code == 403
    assert _JUPYTER_AUTH_CACHE == {}
    assert jupyter_module._jupyter_auth_cache_get(WS_ID, USER_ID) is None


async def test_runtime_start_failure_does_not_populate_cache(monkeypatch) -> None:
    """runtime 启动失败（500）不进缓存。"""
    runtime = _make_runtime_client()
    runtime.get_workspace = AsyncMock(return_value=None)  # 未运行 -> 走 start
    runtime.start_workspace = AsyncMock(
        side_effect=RuntimeClientError(500, {"code": "JUPYTER_START_FAILED"})
    )
    request, response, membership, lock_check, _rt = _setup_mocks(
        monkeypatch, runtime_client=runtime
    )

    with pytest.raises(HTTPException) as excinfo:
        await _call_once(request, response)

    assert excinfo.value.status_code == 500
    assert _JUPYTER_AUTH_CACHE == {}


async def test_cache_keyed_per_user_and_workspace(monkeypatch) -> None:
    """不同 user 共享 workspace 时不串缓存。"""
    request, response, membership, lock_check, runtime = _setup_mocks(monkeypatch)
    await _call_once(request, response)  # USER_A 预热缓存
    assert membership.await_count == 1

    # 换一个 user_id，同一个 workspace -> 缓存 key 不同，必须重新跑完整路径。
    request2, response2 = _make_context(runtime)
    monkeypatch.setattr(
        jupyter_module,
        "verify_jwt_token",
        lambda _token: {"sub": "01USR0000000000000000000B"},
    )
    await _call_once(request2, response2)
    assert membership.await_count == 2
    assert runtime.get_workspace.await_count == 2