ComfyUI-OpenClaw/tests/test_r62_queue_degrade.py

"""
Tests for R62 Queue Submit Degrade and R61 Error Contract.
"""

import os
import sys
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from api.errors import APIError, ErrorCode

# We import submit_prompt inside tests to ensure mocks applied before import if needed,
# although the lazy import inside the function makes it easier.
from services.queue_submit import (
    COMFY_USAGE_SOURCE,
    _build_queue_extra_data,
    submit_prompt,
)


class TestR62QueueDegrade(unittest.IsolatedAsyncioTestCase):

    async def test_dependency_missing(self):
        """Test R62: specific error when aiohttp is missing."""
        # Clean sys.modules to ensure we can control import
        with patch.dict(sys.modules, {"aiohttp": None}):
            with self.assertRaises(APIError) as cm:
                await submit_prompt({"test": "workflow"})

            err = cm.exception
            self.assertEqual(err.code, ErrorCode.DEPENDENCY_UNAVAILABLE)
            self.assertEqual(err.status, 503)
            self.assertIn("required for queue submission", err.message)

    async def test_submit_success(self):
        """Test successful submission when aiohttp is present."""

        mock_response = MagicMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={"prompt_id": "123", "number": 1})

        mock_session_inst = MagicMock()
        mock_session_inst.post.return_value.__aenter__.return_value = mock_response

        mock_session_cls = MagicMock()
        mock_session_cls.return_value.__aenter__.return_value = mock_session_inst

        # Mock aiohttp module
        mock_aiohttp = MagicMock()
        mock_aiohttp.ClientSession = mock_session_cls

        with patch.dict(sys.modules, {"aiohttp": mock_aiohttp}):
            result = await submit_prompt({"test": "workflow"})

            self.assertEqual(result["prompt_id"], "123")
            mock_session_inst.post.assert_called_once()
            sent_payload = mock_session_inst.post.call_args.kwargs["json"]
            self.assertEqual(
                sent_payload["extra_data"]["comfy_usage_source"], COMFY_USAGE_SOURCE
            )
            self.assertIn("tenant_id", sent_payload["extra_data"]["openclaw"])

    def test_queue_extra_data_sets_stable_usage_source_without_prompt_leak(self):
        prompt_text = "private prompt body should not become attribution"

        extra = _build_queue_extra_data(
            {"openclaw": {"trace_id": "trace-1"}, "moltbot": {"trace_id": "trace-1"}},
            tenant_id="tenant-a",
        )

        self.assertEqual(extra["comfy_usage_source"], COMFY_USAGE_SOURCE)
        self.assertEqual(extra["openclaw"]["trace_id"], "trace-1")
        self.assertEqual(extra["openclaw"]["tenant_id"], "tenant-a")
        self.assertEqual(extra["moltbot"]["trace_id"], "trace-1")
        self.assertNotIn(prompt_text, extra["comfy_usage_source"])
        self.assertNotIn("tenant-a", extra["comfy_usage_source"])
        self.assertNotIn("trace-1", extra["comfy_usage_source"])

    async def test_submit_preserves_caller_usage_source_and_openclaw_metadata(self):
        """Caller-provided ComfyUI attribution takes precedence."""

        mock_response = MagicMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={"prompt_id": "123", "number": 1})

        mock_session_inst = MagicMock()
        mock_session_inst.post.return_value.__aenter__.return_value = mock_response

        mock_session_cls = MagicMock()
        mock_session_cls.return_value.__aenter__.return_value = mock_session_inst

        mock_aiohttp = MagicMock()
        mock_aiohttp.ClientSession = mock_session_cls

        caller_extra = {
            "comfy_usage_source": "caller-owned-source",
            "openclaw": {"trace_id": "trace-caller"},
            "moltbot": {"trace_id": "trace-caller"},
        }
        workflow = {
            "1": {
                "class_type": "CLIPTextEncode",
                "inputs": {"text": "private prompt text"},
            }
        }

        with patch.dict(sys.modules, {"aiohttp": mock_aiohttp}):
            result = await submit_prompt(
                workflow,
                extra_data=caller_extra,
                tenant_id="tenant-caller",
            )

        self.assertEqual(result["prompt_id"], "123")
        sent_payload = mock_session_inst.post.call_args.kwargs["json"]
        sent_extra = sent_payload["extra_data"]
        self.assertEqual(sent_extra["comfy_usage_source"], "caller-owned-source")
        self.assertEqual(sent_extra["openclaw"]["trace_id"], "trace-caller")
        self.assertEqual(sent_extra["openclaw"]["tenant_id"], "tenant-caller")
        self.assertEqual(sent_extra["moltbot"]["trace_id"], "trace-caller")
        self.assertNotIn("private prompt text", sent_extra["comfy_usage_source"])
        self.assertEqual(caller_extra["openclaw"], {"trace_id": "trace-caller"})

    async def test_upstream_failure(self):
        """Test standard APIError when upstream returns non-200."""

        mock_response = MagicMock()
        mock_response.status = 500
        mock_response.text = AsyncMock(return_value="Internal Server Error")

        mock_session_inst = MagicMock()
        mock_session_inst.post.return_value.__aenter__.return_value = mock_response

        mock_session_cls = MagicMock()
        mock_session_cls.return_value.__aenter__.return_value = mock_session_inst

        mock_aiohttp = MagicMock()
        mock_aiohttp.ClientSession = mock_session_cls

        with patch.dict(sys.modules, {"aiohttp": mock_aiohttp}):
            with self.assertRaises(APIError) as cm:
                await submit_prompt({"test": "workflow"})

            err = cm.exception
            self.assertEqual(err.code, ErrorCode.QUEUE_SUBMIT_FAILED)
            self.assertEqual(err.status, 502)
            self.assertIn("Queue submission failed: 500", err.message)


if __name__ == "__main__":
    unittest.main()