mizan/backends/mizan-fastapi/tests/test_websocket.py

"""
WebSocket RPC behavior — the genuine capability behind the `websocket` probe.

Proves the `/ws/` route dispatches `@client(websocket=True)` functions through
the SAME `mizan_core.dispatch` core as `POST /call/`: input validation, the
`{result, invalidate, merge}` envelope, `auth=` enforcement, and the
websocket=True gate that rejects HTTP-only functions. The frame protocol matches
mizan-django's Channels consumer (`action:"rpc"` → `{id, ok, data|error}`).
"""

from __future__ import annotations

import pytest
from fastapi import FastAPI
from fastapi.exceptions import RequestValidationError
from fastapi.testclient import TestClient
from pydantic import BaseModel

from mizan_core.client.function import client
from mizan_core.registry import clear_registry, register
from mizan_fastapi import (
    MizanError,
    mizan_exception_handler,
    mizan_validation_handler,
    router as mizan_router,
)


class EchoOut(BaseModel):
    message: str


@pytest.fixture
def app():
    clear_registry()

    @client(websocket=True)
    def ws_echo(request, text: str) -> EchoOut:
        return EchoOut(message=f"ws: {text}")

    @client(websocket=True)
    def ws_add(request, a: int, b: int) -> dict:
        return {"total": a + b}

    @client(websocket=True, affects="user")
    def ws_update(request, user_id: int) -> dict:
        return {"ok": True}

    @client(websocket=True, auth=True)
    def ws_secret(request) -> dict:
        return {"secret": True}

    @client  # HTTP-only — must be rejected over WS
    def http_only(request) -> dict:
        return {"http": True}

    for fn, name in (
        (ws_echo, "ws_echo"), (ws_add, "ws_add"), (ws_update, "ws_update"),
        (ws_secret, "ws_secret"), (http_only, "http_only"),
    ):
        register(fn, name)

    fastapi_app = FastAPI()
    fastapi_app.include_router(mizan_router, prefix="/api/mizan")
    fastapi_app.add_exception_handler(MizanError, mizan_exception_handler)
    fastapi_app.add_exception_handler(RequestValidationError, mizan_validation_handler)
    yield fastapi_app
    clear_registry()


@pytest.fixture
def http(app):
    return TestClient(app)


def test_ws_rpc_dispatches_and_returns_data(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "1", "fn": "ws_echo", "args": {"text": "hi"}})
        frame = ws.receive_json()
    assert frame == {"id": "1", "ok": True, "data": {"message": "ws: hi"}, "invalidate": []}


def test_ws_rpc_validates_input_through_core(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "2", "fn": "ws_add", "args": {"a": "nope", "b": 3}})
        frame = ws.receive_json()
    assert frame["ok"] is False
    assert frame["error"]["code"] == "VALIDATION_ERROR"


def test_ws_rpc_carries_invalidation(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "3", "fn": "ws_update", "args": {"user_id": 5}})
        frame = ws.receive_json()
    assert frame["ok"] is True
    assert "user" in frame["invalidate"]


def test_http_only_function_is_forbidden_over_ws(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "4", "fn": "http_only", "args": {}})
        frame = ws.receive_json()
    assert frame["ok"] is False
    assert frame["error"]["code"] == "FORBIDDEN"


def test_unknown_function_over_ws_is_not_found(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "5", "fn": "ghost", "args": {}})
        frame = ws.receive_json()
    assert frame["ok"] is False
    assert frame["error"]["code"] == "NOT_FOUND"


def test_auth_required_function_rejects_anonymous_over_ws(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "6", "fn": "ws_secret", "args": {}})
        frame = ws.receive_json()
    assert frame["ok"] is False
    assert frame["error"]["code"] == "UNAUTHORIZED"


def test_missing_fn_field_is_bad_request(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "7"})
        frame = ws.receive_json()
    assert frame["ok"] is False
    assert frame["error"]["code"] == "BAD_REQUEST"


def test_unknown_action_errors(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "bogus"})
        frame = ws.receive_json()
    assert "error" in frame


def test_multiple_calls_on_one_connection(http):
    with http.websocket_connect("/api/mizan/ws/") as ws:
        ws.send_json({"action": "rpc", "id": "a", "fn": "ws_echo", "args": {"text": "1"}})
        first = ws.receive_json()
        ws.send_json({"action": "rpc", "id": "b", "fn": "ws_echo", "args": {"text": "2"}})
        second = ws.receive_json()
    assert first["data"]["message"] == "ws: 1"
    assert second["data"]["message"] == "ws: 2"