"""Tests for session queue API behavior with workflow-call queue items.""" import logging import uuid from typing import Any from unittest.mock import MagicMock import pytest from fastapi.testclient import TestClient from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue from invokeai.app.services.shared.graph import Graph, GraphExecutionState from invokeai.app.services.users.users_common import UserCreateRequest from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage from invokeai.backend.util.logging import InvokeAILogger from tests.fixtures.sqlite_database import create_mock_sqlite_database class MockApiDependencies(ApiDependencies): invoker: Invoker def __init__(self, invoker: Invoker) -> None: self.invoker = invoker @pytest.fixture def setup_jwt_secret(): from invokeai.app.services.auth.token_service import set_jwt_secret set_jwt_secret("test-secret-key-for-unit-tests-only-do-not-use-in-production") @pytest.fixture def client(): return TestClient(app) @pytest.fixture def mock_services() -> InvocationServices: from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage from invokeai.app.services.boards.boards_default import BoardService from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ( ClientStatePersistenceSqlite, ) from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage from invokeai.app.services.images.images_default import ImageService from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService from invokeai.app.services.users.users_default import UserService from tests.test_nodes import TestEventService configuration = InvokeAIAppConfig(use_memory_db=True, node_cache_size=0) logger = InvokeAILogger.get_logger() db = create_mock_sqlite_database(configuration, logger) return InvocationServices( board_image_records=SqliteBoardImageRecordStorage(db=db), board_images=None, # type: ignore board_records=SqliteBoardRecordStorage(db=db), boards=BoardService(), bulk_download=BulkDownloadService(), configuration=configuration, events=TestEventService(), image_files=None, # type: ignore image_records=SqliteImageRecordStorage(db=db), images=ImageService(), invocation_cache=MemoryInvocationCache(max_cache_size=0), logger=logging, # type: ignore model_images=None, # type: ignore model_manager=None, # type: ignore download_queue=None, # type: ignore names=None, # type: ignore performance_statistics=InvocationStatsService(), session_processor=None, # type: ignore session_queue=None, # type: ignore urls=None, # type: ignore workflow_records=SqliteWorkflowRecordsStorage(db=db), tensors=None, # type: ignore conditioning=None, # type: ignore style_preset_records=None, # type: ignore style_preset_image_files=None, # type: ignore workflow_thumbnails=None, # type: ignore model_relationship_records=None, # type: ignore model_relationships=None, # type: ignore client_state_persistence=ClientStatePersistenceSqlite(db=db), users=UserService(db), external_generation=None, # type: ignore ) @pytest.fixture def mock_invoker(mock_services: InvocationServices) -> Invoker: invoker = Invoker(services=mock_services) queue = SqliteSessionQueue(db=mock_services.board_records._db) mock_services.session_queue = queue mock_services.session_processor = MagicMock() mock_services.session_processor.get_status.return_value = SessionProcessorStatus( is_started=True, is_processing=False ) queue.start(invoker) return invoker def _create_user(mock_invoker: Invoker, email: str, display_name: str, is_admin: bool = False) -> str: user = mock_invoker.services.users.create( UserCreateRequest(email=email, display_name=display_name, password="TestPass123", is_admin=is_admin) ) return user.user_id def _login(client: TestClient, email: str) -> str: response = client.post("/api/v1/auth/login", json={"email": email, "password": "TestPass123", "remember_me": False}) assert response.status_code == 200 return response.json()["token"] def _auth(token: str) -> dict[str, str]: return {"Authorization": f"Bearer {token}"} @pytest.fixture def enable_multiuser(monkeypatch: Any, mock_invoker: Invoker): mock_invoker.services.configuration.multiuser = True mock_deps = MockApiDependencies(mock_invoker) monkeypatch.setattr("invokeai.app.api.routers.auth.ApiDependencies", mock_deps) monkeypatch.setattr("invokeai.app.api.auth_dependencies.ApiDependencies", mock_deps) monkeypatch.setattr("invokeai.app.api.routers.session_queue.ApiDependencies", mock_deps) yield @pytest.fixture def admin_token(setup_jwt_secret: None, enable_multiuser: Any, mock_invoker: Invoker, client: TestClient): _create_user(mock_invoker, "admin@test.com", "Admin", is_admin=True) return _login(client, "admin@test.com") @pytest.fixture def user1_token(enable_multiuser: Any, mock_invoker: Invoker, client: TestClient, admin_token: str): _create_user(mock_invoker, "user1@test.com", "User One") return _login(client, "user1@test.com") @pytest.fixture def user2_token(enable_multiuser: Any, mock_invoker: Invoker, client: TestClient, admin_token: str): _create_user(mock_invoker, "user2@test.com", "User Two") return _login(client, "user2@test.com") def _insert_queue_item( session_queue: SqliteSessionQueue, *, queue_id: str = "default", user_id: str, status: str, session: GraphExecutionState | None = None, workflow_call_id: str | None = None, parent_item_id: int | None = None, parent_session_id: str | None = None, root_item_id: int | None = None, workflow_call_depth: int | None = None, ) -> int: session = session or GraphExecutionState(graph=Graph()) with session_queue._db.transaction() as cursor: cursor.execute( """--sql INSERT INTO session_queue ( queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id, user_id, status, workflow_call_id, parent_item_id, parent_session_id, root_item_id, workflow_call_depth ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( queue_id, session.model_dump_json(warnings=False), session.id, str(uuid.uuid4()), None, 0, None, None, None, None, user_id, status, workflow_call_id, parent_item_id, parent_session_id, root_item_id, workflow_call_depth, ), ) return cursor.lastrowid def test_get_queue_status_reports_waiting_for_owner( client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") user2 = mock_invoker.services.users.get_by_email("user2@test.com") assert user1 is not None and user2 is not None _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="waiting") _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="pending") response = client.get("/api/v1/queue/default/status", headers=_auth(user1_token)) assert response.status_code == 200 payload = response.json() assert payload["queue"]["waiting"] == 1 assert payload["queue"]["pending"] == 1 assert payload["queue"]["in_progress"] == 0 assert payload["queue"]["total"] == 2 assert payload["queue"]["user_pending"] == 0 assert payload["queue"]["user_in_progress"] == 0 assert payload["queue"]["item_id"] is None def test_get_queue_item_sanitizes_workflow_call_metadata_for_non_owner( client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None item_id = _insert_queue_item( mock_invoker.services.session_queue, user_id=user1.user_id, status="waiting", workflow_call_id="workflow-call-1", parent_item_id=41, parent_session_id="parent-session-1", root_item_id=17, workflow_call_depth=2, ) response = client.get(f"/api/v1/queue/default/i/{item_id}", headers=_auth(user2_token)) assert response.status_code == 200 payload = response.json() assert payload["status"] == "waiting" assert payload["item_id"] == item_id assert payload["user_id"] == "redacted" assert payload["batch_id"] == "redacted" assert payload["session_id"] == "redacted" assert payload.get("workflow_call_id") is None assert payload.get("parent_item_id") is None assert payload.get("parent_session_id") is None assert payload.get("root_item_id") is None assert payload.get("workflow_call_depth") is None def test_retry_items_by_id_normalizes_child_to_root_at_router( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="failed") child_item_id = _insert_queue_item( mock_invoker.services.session_queue, user_id=user1.user_id, status="failed", workflow_call_id="workflow-call-1", parent_item_id=root_item_id, parent_session_id="parent-session-1", root_item_id=root_item_id, workflow_call_depth=1, ) response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[child_item_id]) assert response.status_code == 200 assert response.json()["retried_item_ids"] == [root_item_id] def test_retry_items_by_id_rejects_items_from_other_queue( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None item_id = _insert_queue_item( mock_invoker.services.session_queue, queue_id="secondary", user_id=user1.user_id, status="failed", ) response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[item_id]) assert response.status_code == 404 def test_retry_items_by_id_authorizes_workflow_call_root_owner( client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") user2 = mock_invoker.services.users.get_by_email("user2@test.com") assert user1 is not None and user2 is not None root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="failed") child_item_id = _insert_queue_item( mock_invoker.services.session_queue, user_id=user1.user_id, status="failed", workflow_call_id="workflow-call-1", parent_item_id=root_item_id, parent_session_id="parent-session-1", root_item_id=root_item_id, workflow_call_depth=1, ) response = client.put("/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[child_item_id]) assert response.status_code == 403 assert mock_invoker.services.session_queue.get_queue_item(root_item_id).status == "failed" def test_retry_items_by_id_skips_missing_items_and_retries_valid_items( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="failed") response = client.put( "/api/v1/queue/default/retry_items_by_id", headers=_auth(user1_token), json=[root_item_id + 9999, root_item_id], ) assert response.status_code == 200 assert response.json()["retried_item_ids"] == [root_item_id] def test_cancel_queue_item_cascades_to_waiting_parent_via_router( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None parent_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user1.user_id, status="waiting") child_item_id = _insert_queue_item( mock_invoker.services.session_queue, user_id=user1.user_id, status="pending", workflow_call_id="workflow-call-1", parent_item_id=parent_item_id, parent_session_id="parent-session-1", root_item_id=parent_item_id, workflow_call_depth=1, ) response = client.put(f"/api/v1/queue/default/i/{child_item_id}/cancel", headers=_auth(user1_token)) assert response.status_code == 200 assert response.json()["status"] == "canceled" assert mock_invoker.services.session_queue.get_queue_item(parent_item_id).status == "canceled" assert mock_invoker.services.session_queue.get_queue_item(child_item_id).status == "canceled" def test_cancel_queue_item_rejects_item_from_other_queue( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None item_id = _insert_queue_item( mock_invoker.services.session_queue, queue_id="secondary", user_id=user1.user_id, status="pending", ) response = client.put(f"/api/v1/queue/default/i/{item_id}/cancel", headers=_auth(user1_token)) assert response.status_code == 404 def test_delete_queue_item_rejects_item_from_other_queue( client: TestClient, mock_invoker: Invoker, user1_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") assert user1 is not None item_id = _insert_queue_item( mock_invoker.services.session_queue, queue_id="secondary", user_id=user1.user_id, status="failed", ) response = client.delete(f"/api/v1/queue/default/i/{item_id}", headers=_auth(user1_token)) assert response.status_code == 404 def test_delete_queue_item_authorizes_workflow_call_root_owner( client: TestClient, mock_invoker: Invoker, user1_token: str, user2_token: str ) -> None: user1 = mock_invoker.services.users.get_by_email("user1@test.com") user2 = mock_invoker.services.users.get_by_email("user2@test.com") assert user1 is not None and user2 is not None root_item_id = _insert_queue_item(mock_invoker.services.session_queue, user_id=user2.user_id, status="failed") child_item_id = _insert_queue_item( mock_invoker.services.session_queue, user_id=user1.user_id, status="failed", workflow_call_id="workflow-call-1", parent_item_id=root_item_id, parent_session_id="parent-session-1", root_item_id=root_item_id, workflow_call_depth=1, ) response = client.delete(f"/api/v1/queue/default/i/{child_item_id}", headers=_auth(user1_token)) assert response.status_code == 403 assert mock_invoker.services.session_queue.get_queue_item(root_item_id).status == "failed" assert mock_invoker.services.session_queue.get_queue_item(child_item_id).status == "failed"