kantan-kanto-ComfyUI-MultiModal-Prompt-Nodes/tests/test_qwen_api_models.py

import re
import unittest
from pathlib import Path
from types import SimpleNamespace

from qwen_api_models import (
    QWEN_API_MODEL_IDS,
    QWEN_API_MODELS,
    QWEN_API_MULTIMODAL_ENDPOINT_MODEL_IDS,
    QWEN_API_RETIRED_MODEL_IDS,
    QWEN_API_VISION_MODEL_IDS,
    normalize_api_model_name,
)
from qwen_api_utils import call_qwen_api, extract_qwen_api_response_text


class FakeEndpoint:
    def __init__(self):
        self.calls = []

    def call(self, **kwargs):
        self.calls.append(kwargs)
        return SimpleNamespace(status_code=200)


class QwenAPIModelTests(unittest.TestCase):
    def test_new_models_are_registered_with_expected_capabilities(self):
        self.assertIn("qwen3.8-max", QWEN_API_MODEL_IDS)
        self.assertIn("qwen3.8-2.4t-a95b", QWEN_API_MODEL_IDS)
        self.assertIn("qwen3.7-flash", QWEN_API_MODEL_IDS)
        self.assertIn("qwen3.7-flash-2026-07-15", QWEN_API_MODEL_IDS)
        self.assertIn("qwen3.7-max-2026-06-08", QWEN_API_MODEL_IDS)

        self.assertIn("qwen3.8-max", QWEN_API_VISION_MODEL_IDS)
        self.assertNotIn("qwen3.8-2.4t-a95b", QWEN_API_VISION_MODEL_IDS)
        self.assertIn("qwen3.7-max-2026-06-08", QWEN_API_VISION_MODEL_IDS)
        self.assertNotIn("qwen3.7-max", QWEN_API_VISION_MODEL_IDS)
        self.assertNotIn("qwen3.7-max-2026-05-20", QWEN_API_VISION_MODEL_IDS)

    def test_retired_models_are_not_selectable(self):
        self.assertTrue(QWEN_API_RETIRED_MODEL_IDS.isdisjoint(QWEN_API_MODEL_IDS))

    def test_current_retirement_labels_are_normalized(self):
        labels = [
            label for label in QWEN_API_MODELS if label.startswith("qwen-vl-max ")
        ]
        self.assertEqual(len(labels), 1)
        self.assertIn("offline scheduled 2026-10-10", labels[0])
        self.assertIn("use qwen3.7-plus", labels[0])
        self.assertEqual(normalize_api_model_name(labels[0]), "qwen-vl-max")

        self.assertIn("qwen-plus", QWEN_API_MODELS)
        self.assertIn("qwen-max", QWEN_API_MODELS)

    def test_qwen38_uses_multimodal_endpoint_for_text(self):
        generation = FakeEndpoint()
        multimodal = FakeEndpoint()
        dashscope = SimpleNamespace(
            Generation=generation,
            MultiModalConversation=multimodal,
        )

        call_qwen_api(
            dashscope,
            api_key="test-key",
            model="qwen3.8-max",
            messages=[{"role": "user", "content": "hello"}],
        )

        self.assertEqual(generation.calls, [])
        self.assertEqual(len(multimodal.calls), 1)
        self.assertEqual(
            multimodal.calls[0]["messages"][0]["content"],
            [{"text": "hello"}],
        )
        self.assertFalse(multimodal.calls[0]["enable_thinking"])

    def test_qwen38_multimodal_forwards_json_response_format(self):
        generation = FakeEndpoint()
        multimodal = FakeEndpoint()
        dashscope = SimpleNamespace(
            Generation=generation,
            MultiModalConversation=multimodal,
        )
        response_format = {"type": "json_object"}

        call_qwen_api(
            dashscope,
            api_key="test-key",
            model="qwen3.8-max",
            messages=[
                {"role": "user", "content": "Return the result as JSON."}
            ],
            response_format=response_format,
            force_multimodal=True,
        )

        self.assertEqual(generation.calls, [])
        self.assertEqual(
            multimodal.calls[0]["response_format"], response_format
        )
        self.assertFalse(multimodal.calls[0]["enable_thinking"])

    def test_qwen38_open_model_uses_generation_endpoint(self):
        generation = FakeEndpoint()
        multimodal = FakeEndpoint()
        dashscope = SimpleNamespace(
            Generation=generation,
            MultiModalConversation=multimodal,
        )

        call_qwen_api(
            dashscope,
            api_key="test-key",
            model="qwen3.8-2.4t-a95b",
            messages=[{"role": "user", "content": "hello"}],
        )

        self.assertEqual(len(generation.calls), 1)
        self.assertEqual(multimodal.calls, [])
        self.assertFalse(generation.calls[0]["enable_thinking"])

    def test_vision_models_derive_multimodal_endpoint_membership(self):
        self.assertTrue(
            QWEN_API_VISION_MODEL_IDS.issubset(
                QWEN_API_MULTIMODAL_ENDPOINT_MODEL_IDS
            )
        )

    def test_frontend_vision_filter_matches_api_model_registry(self):
        script_path = (
            Path(__file__).resolve().parents[1] / "web" / "mmproj_filter.js"
        )
        script = script_path.read_text(encoding="utf-8")
        match = re.search(
            r"const VISION_API_MODEL_IDS = new Set\(\[(.*?)\]\);",
            script,
            flags=re.DOTALL,
        )

        self.assertIsNotNone(match)
        frontend_model_ids = set(re.findall(r'"([^"]+)"', match.group(1)))
        self.assertEqual(frontend_model_ids, QWEN_API_VISION_MODEL_IDS)

    def test_extracts_text_from_supported_response_shapes(self):
        responses = [
            SimpleNamespace(
                output=SimpleNamespace(
                    choices=[
                        SimpleNamespace(
                            message=SimpleNamespace(content="plain text")
                        )
                    ]
                )
            ),
            {
                "output": {
                    "choices": [
                        {
                            "message": {
                                "content": [
                                    {"text": "first "},
                                    {"text": "second"},
                                ]
                            }
                        }
                    ]
                }
            },
            SimpleNamespace(output=SimpleNamespace(text="output text")),
        ]

        self.assertEqual(extract_qwen_api_response_text(responses[0]), "plain text")
        self.assertEqual(extract_qwen_api_response_text(responses[1]), "first second")
        self.assertEqual(extract_qwen_api_response_text(responses[2]), "output text")


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