InvokeAI/tests/app/routers/test_client_state_multiuser.py

"""Tests for multiuser client state functionality."""

from typing import Any

import pytest
from fastapi import status
from fastapi.testclient import TestClient

from invokeai.app.api.dependencies import ApiDependencies
from invokeai.app.api_app import app
from invokeai.app.services.invoker import Invoker
from invokeai.app.services.users.users_common import UserCreateRequest


@pytest.fixture
def client():
    """Create a test client."""
    return TestClient(app)


class MockApiDependencies(ApiDependencies):
    """Mock API dependencies for testing."""

    invoker: Invoker

    def __init__(self, invoker: Invoker) -> None:
        self.invoker = invoker


def setup_test_user(
    mock_invoker: Invoker, email: str, display_name: str, password: str = "TestPass123", is_admin: bool = False
) -> str:
    """Helper to create a test user and return user_id."""
    user_service = mock_invoker.services.users
    user_data = UserCreateRequest(
        email=email,
        display_name=display_name,
        password=password,
        is_admin=is_admin,
    )
    user = user_service.create(user_data)
    return user.user_id


def get_user_token(client: TestClient, email: str, password: str = "TestPass123") -> str:
    """Helper to login and get a user token."""
    response = client.post(
        "/api/v1/auth/login",
        json={
            "email": email,
            "password": password,
            "remember_me": False,
        },
    )
    assert response.status_code == 200
    return response.json()["token"]


@pytest.fixture
def admin_token(monkeypatch: Any, mock_invoker: Invoker, client: TestClient):
    """Get an admin token for testing."""
    # Enable multiuser mode for auth endpoints
    mock_invoker.services.configuration.multiuser = True

    # Mock ApiDependencies for auth and client_state routers
    monkeypatch.setattr("invokeai.app.api.routers.auth.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Create admin user
    setup_test_user(mock_invoker, "admin@test.com", "Admin User", is_admin=True)

    return get_user_token(client, "admin@test.com")


@pytest.fixture
def user1_token(monkeypatch: Any, mock_invoker: Invoker, client: TestClient, admin_token: str):
    """Get a token for test user 1."""
    # Create a regular user
    setup_test_user(mock_invoker, "user1@test.com", "User One", is_admin=False)

    return get_user_token(client, "user1@test.com")


@pytest.fixture
def user2_token(monkeypatch: Any, mock_invoker: Invoker, client: TestClient, admin_token: str):
    """Get a token for test user 2."""
    # Create another regular user
    setup_test_user(mock_invoker, "user2@test.com", "User Two", is_admin=False)

    return get_user_token(client, "user2@test.com")


def test_get_client_state_without_auth_uses_system_user(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that getting client state without authentication uses the system user."""
    # Mock ApiDependencies
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Set a value for the system user directly
    mock_invoker.services.client_state_persistence.set_by_key("system", "test_key", "system_value")

    # Get without authentication - should return system user's value
    response = client.get("/api/v1/client_state/default/get_by_key?key=test_key")
    assert response.status_code == status.HTTP_200_OK
    assert response.json() == "system_value"


def test_set_client_state_without_auth_uses_system_user(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that setting client state without authentication uses the system user."""
    # Mock ApiDependencies
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Set without authentication - should set for system user
    response = client.post(
        "/api/v1/client_state/default/set_by_key?key=test_key",
        json="unauthenticated_value",
    )
    assert response.status_code == status.HTTP_200_OK

    # Verify it was set for system user
    value = mock_invoker.services.client_state_persistence.get_by_key("system", "test_key")
    assert value == "unauthenticated_value"


def test_delete_client_state_without_auth_uses_system_user(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that deleting client state without authentication uses the system user."""
    # Mock ApiDependencies
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Set a value for system user
    mock_invoker.services.client_state_persistence.set_by_key("system", "test_key", "system_value")

    # Delete without authentication - should delete system user's data
    response = client.post("/api/v1/client_state/default/delete")
    assert response.status_code == status.HTTP_200_OK

    # Verify it was deleted for system user
    value = mock_invoker.services.client_state_persistence.get_by_key("system", "test_key")
    assert value is None


def test_set_and_get_client_state(client: TestClient, admin_token: str):
    """Test that authenticated users can set and get their client state."""
    # Set a value
    set_response = client.post(
        "/api/v1/client_state/default/set_by_key?key=test_key",
        json="test_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert set_response.status_code == status.HTTP_200_OK
    assert set_response.json() == "test_value"

    # Get the value back
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=test_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.status_code == status.HTTP_200_OK
    assert get_response.json() == "test_value"


def test_client_state_isolation_between_users(client: TestClient, user1_token: str, user2_token: str):
    """Test that client state is isolated between different users."""
    # User 1 sets a value
    user1_set_response = client.post(
        "/api/v1/client_state/default/set_by_key?key=shared_key",
        json="user1_value",
        headers={"Authorization": f"Bearer {user1_token}"},
    )
    assert user1_set_response.status_code == status.HTTP_200_OK

    # User 2 sets a different value for the same key
    user2_set_response = client.post(
        "/api/v1/client_state/default/set_by_key?key=shared_key",
        json="user2_value",
        headers={"Authorization": f"Bearer {user2_token}"},
    )
    assert user2_set_response.status_code == status.HTTP_200_OK

    # User 1 should still see their own value
    user1_get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=shared_key",
        headers={"Authorization": f"Bearer {user1_token}"},
    )
    assert user1_get_response.status_code == status.HTTP_200_OK
    assert user1_get_response.json() == "user1_value"

    # User 2 should see their own value
    user2_get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=shared_key",
        headers={"Authorization": f"Bearer {user2_token}"},
    )
    assert user2_get_response.status_code == status.HTTP_200_OK
    assert user2_get_response.json() == "user2_value"


def test_get_nonexistent_key_returns_null(client: TestClient, admin_token: str):
    """Test that getting a nonexistent key returns null."""
    response = client.get(
        "/api/v1/client_state/default/get_by_key?key=nonexistent_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    assert response.json() is None


def test_delete_client_state(client: TestClient, admin_token: str):
    """Test that users can delete their own client state."""
    # Set some values
    client.post(
        "/api/v1/client_state/default/set_by_key?key=key1",
        json="value1",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    client.post(
        "/api/v1/client_state/default/set_by_key?key=key2",
        json="value2",
        headers={"Authorization": f"Bearer {admin_token}"},
    )

    # Verify values exist
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=key1",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.json() == "value1"

    # Delete all client state
    delete_response = client.post(
        "/api/v1/client_state/default/delete",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert delete_response.status_code == status.HTTP_200_OK

    # Verify values are gone
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=key1",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.json() is None

    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=key2",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.json() is None


def test_update_existing_key(client: TestClient, admin_token: str):
    """Test that updating an existing key works correctly."""
    # Set initial value
    client.post(
        "/api/v1/client_state/default/set_by_key?key=update_key",
        json="initial_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )

    # Update the value
    update_response = client.post(
        "/api/v1/client_state/default/set_by_key?key=update_key",
        json="updated_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert update_response.status_code == status.HTTP_200_OK

    # Verify the updated value
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=update_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.status_code == status.HTTP_200_OK
    assert get_response.json() == "updated_value"


def test_complex_json_values(client: TestClient, admin_token: str):
    """Test that complex JSON values can be stored and retrieved."""
    import json

    complex_dict = {"params": {"model": "test-model", "steps": 50}, "prompt": "a beautiful landscape"}
    complex_value = json.dumps(complex_dict)

    # Set complex value
    set_response = client.post(
        "/api/v1/client_state/default/set_by_key?key=complex_key",
        json=complex_value,
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert set_response.status_code == status.HTTP_200_OK

    # Get it back
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=complex_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.status_code == status.HTTP_200_OK
    assert get_response.json() == complex_value


def test_get_keys_by_prefix_without_auth(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that keys can be retrieved by prefix without authentication."""
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Set several keys with a common prefix directly
    for i in range(3):
        mock_invoker.services.client_state_persistence.set_by_key("system", f"canvas_snapshot:snap{i}", f"value{i}")
    mock_invoker.services.client_state_persistence.set_by_key("system", "other_key", "other_value")

    # Get keys by prefix
    response = client.get("/api/v1/client_state/default/get_keys_by_prefix?prefix=canvas_snapshot:")
    assert response.status_code == status.HTTP_200_OK
    keys = response.json()
    assert len(keys) == 3
    assert "canvas_snapshot:snap0" in keys
    assert "canvas_snapshot:snap1" in keys
    assert "canvas_snapshot:snap2" in keys
    assert "other_key" not in keys


def test_get_keys_by_prefix_empty_without_auth(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that an empty list is returned when no keys match the prefix."""
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    response = client.get("/api/v1/client_state/default/get_keys_by_prefix?prefix=nonexistent_prefix:")
    assert response.status_code == status.HTTP_200_OK
    assert response.json() == []


def test_delete_by_key_without_auth(client: TestClient, monkeypatch, mock_invoker: Invoker):
    """Test that a specific key can be deleted without affecting other keys."""
    monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", MockApiDependencies(mock_invoker))
    monkeypatch.setattr("invokeai.app.api.routers.client_state.ApiDependencies", MockApiDependencies(mock_invoker))

    # Set two keys directly
    mock_invoker.services.client_state_persistence.set_by_key("system", "keep_key", "keep_value")
    mock_invoker.services.client_state_persistence.set_by_key("system", "delete_key", "delete_value")

    # Delete only one key via endpoint
    delete_response = client.post("/api/v1/client_state/default/delete_by_key?key=delete_key")
    assert delete_response.status_code == status.HTTP_200_OK

    # Verify deleted key is gone
    value = mock_invoker.services.client_state_persistence.get_by_key("system", "delete_key")
    assert value is None

    # Verify other key still exists
    value = mock_invoker.services.client_state_persistence.get_by_key("system", "keep_key")
    assert value == "keep_value"


def test_get_keys_by_prefix(client: TestClient, admin_token: str):
    """Test that keys can be retrieved by prefix with authentication."""
    # Set several keys with a common prefix
    for i in range(3):
        client.post(
            f"/api/v1/client_state/default/set_by_key?key=canvas_snapshot:snap{i}",
            json=f"value{i}",
            headers={"Authorization": f"Bearer {admin_token}"},
        )
    # Set a key without the prefix
    client.post(
        "/api/v1/client_state/default/set_by_key?key=other_key",
        json="other_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )

    # Get keys by prefix
    response = client.get(
        "/api/v1/client_state/default/get_keys_by_prefix?prefix=canvas_snapshot:",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert response.status_code == status.HTTP_200_OK
    keys = response.json()
    assert len(keys) == 3
    assert "canvas_snapshot:snap0" in keys
    assert "canvas_snapshot:snap1" in keys
    assert "canvas_snapshot:snap2" in keys
    assert "other_key" not in keys


def test_delete_by_key(client: TestClient, admin_token: str):
    """Test that a specific key can be deleted without affecting other keys."""
    # Set two keys
    client.post(
        "/api/v1/client_state/default/set_by_key?key=keep_key",
        json="keep_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    client.post(
        "/api/v1/client_state/default/set_by_key?key=delete_key",
        json="delete_value",
        headers={"Authorization": f"Bearer {admin_token}"},
    )

    # Delete only one key
    delete_response = client.post(
        "/api/v1/client_state/default/delete_by_key?key=delete_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert delete_response.status_code == status.HTTP_200_OK

    # Verify deleted key is gone
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=delete_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.json() is None

    # Verify other key still exists
    get_response = client.get(
        "/api/v1/client_state/default/get_by_key?key=keep_key",
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert get_response.json() == "keep_value"


def test_get_keys_by_prefix_isolation_between_users(client: TestClient, user1_token: str, user2_token: str):
    """Test that get_keys_by_prefix is isolated between users."""
    # User 1 sets keys
    client.post(
        "/api/v1/client_state/default/set_by_key?key=snapshot:u1",
        json="user1_data",
        headers={"Authorization": f"Bearer {user1_token}"},
    )

    # User 2 sets keys
    client.post(
        "/api/v1/client_state/default/set_by_key?key=snapshot:u2",
        json="user2_data",
        headers={"Authorization": f"Bearer {user2_token}"},
    )

    # User 1 should only see their own keys
    response = client.get(
        "/api/v1/client_state/default/get_keys_by_prefix?prefix=snapshot:",
        headers={"Authorization": f"Bearer {user1_token}"},
    )
    keys = response.json()
    assert "snapshot:u1" in keys
    assert "snapshot:u2" not in keys