invoke-ai--invokeai/tests/app/test_workflow_socketio.py

from types import SimpleNamespace
from unittest.mock import ANY, AsyncMock

import pytest
from fastapi import FastAPI

from invokeai.app.api.sockets import SocketIO


@pytest.fixture
def anyio_backend() -> str:
    return "asyncio"


def _patch_multiuser_context(monkeypatch: pytest.MonkeyPatch, *, user_id: str, is_admin: bool) -> None:
    user = SimpleNamespace(user_id=user_id, is_active=True)
    invoker = SimpleNamespace(
        services=SimpleNamespace(
            configuration=SimpleNamespace(multiuser=True),
            users=SimpleNamespace(get=lambda candidate_user_id: user if candidate_user_id == user_id else None),
        )
    )
    monkeypatch.setattr("invokeai.app.api.dependencies.ApiDependencies", SimpleNamespace(invoker=invoker))
    monkeypatch.setattr(
        "invokeai.app.api.sockets.verify_token",
        lambda token: SimpleNamespace(user_id=user_id, is_admin=is_admin) if token == "valid-token" else None,
    )


def _patch_single_user_context(monkeypatch: pytest.MonkeyPatch) -> None:
    invoker = SimpleNamespace(services=SimpleNamespace(configuration=SimpleNamespace(multiuser=False)))
    monkeypatch.setattr("invokeai.app.api.dependencies.ApiDependencies", SimpleNamespace(invoker=invoker))


@pytest.mark.anyio
async def test_authenticated_user_joins_workflow_rooms_on_connect(monkeypatch: pytest.MonkeyPatch) -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.enter_room = AsyncMock()
    _patch_multiuser_context(monkeypatch, user_id="user-1", is_admin=False)

    accepted = await socketio._handle_connect("sid-1", {}, {"token": "valid-token"})

    assert accepted is True
    socketio._sio.enter_room.assert_any_call("sid-1", "user:user-1")
    socketio._sio.enter_room.assert_any_call("sid-1", "workflows:shared")


@pytest.mark.anyio
async def test_admin_joins_admin_room_on_connect(monkeypatch: pytest.MonkeyPatch) -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.enter_room = AsyncMock()
    _patch_multiuser_context(monkeypatch, user_id="admin-1", is_admin=True)

    accepted = await socketio._handle_connect("sid-1", {}, {"token": "valid-token"})

    assert accepted is True
    socketio._sio.enter_room.assert_any_call("sid-1", "user:admin-1")
    socketio._sio.enter_room.assert_any_call("sid-1", "workflows:shared")
    socketio._sio.enter_room.assert_any_call("sid-1", "admin")


@pytest.mark.anyio
async def test_single_user_socket_joins_workflow_rooms_on_connect(monkeypatch: pytest.MonkeyPatch) -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.enter_room = AsyncMock()
    _patch_single_user_context(monkeypatch)

    accepted = await socketio._handle_connect("sid-1", {}, None)

    assert accepted is True
    socketio._sio.enter_room.assert_any_call("sid-1", "user:system")
    socketio._sio.enter_room.assert_any_call("sid-1", "workflows:shared")
    socketio._sio.enter_room.assert_any_call("sid-1", "admin")


@pytest.mark.anyio
async def test_private_workflow_event_is_emitted_only_to_owner_and_admin() -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.emit = AsyncMock()

    event_payload = SimpleNamespace(
        __event_name__="workflow_created",
        workflow_id="wf-1",
        user_id="owner-1",
        is_public=False,
        model_dump=lambda mode="json": {"workflow_id": "wf-1", "user_id": "owner-1", "is_public": False},
    )

    await socketio._handle_workflow_event(("workflow_created", event_payload))

    socketio._sio.emit.assert_any_call(
        event="workflow_created",
        data={"workflow_id": "wf-1", "user_id": "owner-1", "is_public": False},
        room="user:owner-1",
    )
    socketio._sio.emit.assert_any_call(
        event="workflow_created",
        data={"workflow_id": "wf-1", "user_id": "owner-1", "is_public": False},
        room="admin",
    )
    assert socketio._sio.emit.await_count == 2


@pytest.mark.anyio
async def test_single_user_workflow_event_is_emitted_once_to_admin_room(monkeypatch: pytest.MonkeyPatch) -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.emit = AsyncMock()
    _patch_single_user_context(monkeypatch)

    event_payload = SimpleNamespace(
        __event_name__="workflow_created",
        workflow_id="wf-1",
        user_id="system",
        is_public=False,
        model_dump=lambda mode="json": {"workflow_id": "wf-1", "user_id": "system", "is_public": False},
    )

    await socketio._handle_workflow_event(("workflow_created", event_payload))

    socketio._sio.emit.assert_awaited_once_with(
        event="workflow_created",
        data={"workflow_id": "wf-1", "user_id": "system", "is_public": False},
        room="admin",
    )


@pytest.mark.anyio
async def test_shared_workflow_event_is_emitted_to_shared_room() -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.emit = AsyncMock()

    event_payload = SimpleNamespace(
        __event_name__="workflow_updated",
        workflow_id="wf-1",
        user_id="owner-1",
        old_is_public=False,
        new_is_public=True,
        model_dump=lambda mode="json": {
            "workflow_id": "wf-1",
            "user_id": "owner-1",
            "old_is_public": False,
            "new_is_public": True,
        },
    )

    await socketio._handle_workflow_event(("workflow_updated", event_payload))

    socketio._sio.emit.assert_any_call(
        event="workflow_updated",
        data={"workflow_id": "wf-1", "user_id": "owner-1", "old_is_public": False, "new_is_public": True},
        room="workflows:shared",
    )


@pytest.mark.anyio
async def test_shared_to_private_transition_emits_access_revoked_to_shared_room() -> None:
    socketio = SocketIO(FastAPI())
    socketio._sio.emit = AsyncMock()

    event_payload = SimpleNamespace(
        __event_name__="workflow_updated",
        workflow_id="wf-1",
        user_id="owner-1",
        old_is_public=True,
        new_is_public=False,
        model_dump=lambda mode="json": {
            "workflow_id": "wf-1",
            "user_id": "owner-1",
            "old_is_public": True,
            "new_is_public": False,
        },
    )

    await socketio._handle_workflow_event(("workflow_updated", event_payload))

    socketio._sio.emit.assert_any_call(
        event="workflow_access_revoked",
        data={"workflow_id": "wf-1", "user_id": "owner-1", "timestamp": ANY},
        room="workflows:shared",
    )