import importlib import sys import unittest from types import SimpleNamespace from unittest import mock try: from . import test_path except ImportError: import test_path def _completion(text): return SimpleNamespace( choices=[SimpleNamespace(message=SimpleNamespace(content=text))] ) def _run_node(**overrides): from ComfyUI_LLM_API.node import LLMAPIChatNode kwargs = { "provider_profile": "custom_text_image", "api_baseurl": "https://example.com/v1", "api_key": "sk-secret", "api_key_env": "", "model": "demo", "system_prompt": "system", "prompt": "hello", "temperature": 0.6, "execution_seed": 123, "filter_thinking": True, "max_images": 4, "image_preset": "balanced", "custom_image_max_mb": 1.5, "custom_image_max_side": 1280, } kwargs.update(overrides) return LLMAPIChatNode().run_llmapi(**kwargs) class NodeContractTests(unittest.TestCase): def test_package_exports_only_new_node_mapping(self): import ComfyUI_LLM_API self.assertEqual( list(ComfyUI_LLM_API.NODE_CLASS_MAPPINGS), ["LLM_API_CHAT_NODE"], ) self.assertEqual( list(ComfyUI_LLM_API.NODE_DISPLAY_NAME_MAPPINGS), ["LLM_API_CHAT_NODE"], ) self.assertEqual( ComfyUI_LLM_API.NODE_DISPLAY_NAME_MAPPINGS["LLM_API_CHAT_NODE"], "LLM API Chat", ) def test_input_contract_matches_refactored_chat_node(self): from ComfyUI_LLM_API.image_presets import IMAGE_PRESET_NAMES from ComfyUI_LLM_API.node import LLMAPIChatNode from ComfyUI_LLM_API.provider_profiles import PROVIDER_PROFILE_NAMES inputs = LLMAPIChatNode.INPUT_TYPES() required = inputs["required"] self.assertEqual( list(required), [ "provider_profile", "api_baseurl", "api_key", "api_key_env", "model", "system_prompt", "prompt", "temperature", "execution_seed", "filter_thinking", "max_images", "image_preset", "custom_image_max_mb", "custom_image_max_side", ], ) self.assertEqual(required["provider_profile"][0], PROVIDER_PROFILE_NAMES) self.assertEqual( required["provider_profile"][1]["default"], PROVIDER_PROFILE_NAMES[0], ) self.assertEqual(required["api_baseurl"][1]["default"], "") self.assertEqual(required["api_key_env"][1]["default"], "") self.assertTrue(required["api_key_env"][1]["advanced"]) self.assertEqual(required["image_preset"][0], IMAGE_PRESET_NAMES) self.assertEqual(required["image_preset"][1]["default"], "balanced") self.assertEqual(required["custom_image_max_mb"][1]["default"], 1.5) self.assertEqual(required["custom_image_max_side"][1]["default"], 1280) self.assertIn("images", inputs["optional"]) self.assertNotIn("role", inputs["required"]) self.assertNotIn("seed", inputs["required"]) self.assertNotIn("video", inputs["optional"]) self.assertEqual(LLMAPIChatNode.RETURN_TYPES, ("STRING", "STRING")) self.assertEqual(LLMAPIChatNode.RETURN_NAMES, ("response", "error")) def test_package_import_does_not_import_openai(self): removed_modules = { name: sys.modules.pop(name) for name in list(sys.modules) if name == "ComfyUI_LLM_API" or name.startswith("ComfyUI_LLM_API.") } real_import = __import__ def fail_openai_import(name, *args, **kwargs): if name == "openai" or name.startswith("openai."): raise ImportError("openai intentionally unavailable") return real_import(name, *args, **kwargs) try: with mock.patch("builtins.__import__", side_effect=fail_openai_import): package = importlib.import_module("ComfyUI_LLM_API") self.assertIn("LLM_API_CHAT_NODE", package.NODE_CLASS_MAPPINGS) finally: for name in list(sys.modules): if name == "ComfyUI_LLM_API" or name.startswith("ComfyUI_LLM_API."): sys.modules.pop(name) sys.modules.update(removed_modules) def test_validate_inputs_rejects_missing_model(self): from ComfyUI_LLM_API.node import LLMAPIChatNode result = LLMAPIChatNode.VALIDATE_INPUTS( provider_profile="custom_text_image", api_baseurl="https://example.com/v1", model="", temperature=0.6, max_images=4, image_preset="balanced", custom_image_max_mb=1.5, custom_image_max_side=1280, ) self.assertIsInstance(result, str) self.assertIn("model", result.lower()) def test_validate_inputs_rejects_image_limits_above_declared_bounds(self): from ComfyUI_LLM_API.node import LLMAPIChatNode too_many_bytes = LLMAPIChatNode.VALIDATE_INPUTS( provider_profile="custom_text_image", api_baseurl="https://example.com/v1", model="demo", temperature=0.6, max_images=4, image_preset="custom", custom_image_max_mb=32.1, custom_image_max_side=1280, ) too_large_side = LLMAPIChatNode.VALIDATE_INPUTS( provider_profile="custom_text_image", api_baseurl="https://example.com/v1", model="demo", temperature=0.6, max_images=4, image_preset="custom", custom_image_max_mb=1.5, custom_image_max_side=8_193, ) self.assertIn("custom_image_max_mb", too_many_bytes) self.assertIn("custom_image_max_side", too_large_side) def test_validate_inputs_returns_errors_for_non_numeric_runtime_inputs(self): from ComfyUI_LLM_API.node import LLMAPIChatNode cases = ( ("temperature", "warm", 4, "temperature"), ("max_images", 0.6, "many", "max_images"), ) for field, temperature, max_images, expected in cases: with self.subTest(field=field): result = LLMAPIChatNode.VALIDATE_INPUTS( provider_profile="custom_text_image", api_baseurl="https://example.com/v1", model="demo", temperature=temperature, max_images=max_images, image_preset="balanced", custom_image_max_mb=1.5, custom_image_max_side=1280, ) self.assertIsInstance(result, str) self.assertIn(expected, result) def test_validate_inputs_returns_errors_for_bad_custom_image_values(self): from ComfyUI_LLM_API.node import LLMAPIChatNode result = LLMAPIChatNode.VALIDATE_INPUTS( provider_profile="custom_text_image", api_baseurl="https://example.com/v1", model="demo", temperature=0.6, max_images=4, image_preset="custom", custom_image_max_mb=None, custom_image_max_side=1280, ) self.assertIsInstance(result, str) self.assertIn("custom_image_max_mb", result) def test_type_errors_are_not_converted_to_llm_output(self): with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.build_chat_messages", side_effect=TypeError("programming bug"), ): with self.assertRaises(TypeError): _run_node() def test_runtime_errors_are_not_converted_to_llm_output(self): with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.build_chat_messages", side_effect=RuntimeError("programming bug"), ): with self.assertRaises(RuntimeError): _run_node() def test_successful_fake_completion_returns_response_and_empty_error(self): with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", return_value=_completion("hello"), ) as create_completion: response = _run_node() self.assertEqual(response, ("hello", "")) _, kwargs = create_completion.call_args self.assertNotIn("execution_seed", kwargs) def test_create_openai_client_user_error_returns_error_output(self): from ComfyUI_LLM_API.node import LLMAPIUserError with mock.patch( "ComfyUI_LLM_API.node.create_openai_client", side_effect=LLMAPIUserError( "OpenAI Python package is not installed. Run pip install -r requirements.txt." ), ): with mock.patch("ComfyUI_LLM_API.node.logger.warning"): response = _run_node() self.assertEqual(response[0], "") self.assertTrue(response[1].startswith("LLM API Error: ")) self.assertIn("requirements.txt", response[1]) def test_api_key_env_overrides_plain_api_key(self): with mock.patch.dict("os.environ", {"LLM_API_TEST_KEY": "env-secret"}, clear=False): with mock.patch( "ComfyUI_LLM_API.node.create_openai_client", return_value=object(), ) as create_client: with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", return_value=_completion("ok"), ): response = _run_node( api_key="plain-secret", api_key_env="LLM_API_TEST_KEY", ) self.assertEqual(response, ("ok", "")) self.assertEqual(create_client.call_args.kwargs["api_key"], "env-secret") def test_missing_api_key_env_returns_error_before_client(self): with mock.patch.dict("os.environ", {}, clear=True): with mock.patch("ComfyUI_LLM_API.node.create_openai_client") as create_client: with mock.patch("ComfyUI_LLM_API.node.logger.warning"): response = _run_node(api_key_env="LLM_API_MISSING_KEY") self.assertEqual(response[0], "") self.assertIn("Environment variable LLM_API_MISSING_KEY is not set", response[1]) create_client.assert_not_called() def test_api_key_env_value_is_redacted_from_errors(self): with mock.patch.dict("os.environ", {"LLM_API_TEST_KEY": "env-secret"}, clear=False): with mock.patch("ComfyUI_LLM_API.node.logger.warning") as warning: with mock.patch( "ComfyUI_LLM_API.node.create_openai_client", side_effect=ValueError("bad env-secret"), ): response = _run_node(api_key_env="LLM_API_TEST_KEY") self.assertEqual(response[0], "") self.assertNotIn("env-secret", response[1]) self.assertIn("[redacted]", response[1]) self.assertNotIn( "env-secret", " ".join(str(arg) for arg in warning.call_args.args), ) def test_text_only_profile_with_images_returns_error_before_encoding_or_client(self): with mock.patch( "ComfyUI_LLM_API.node.encode_comfy_image_batch_to_data_urls" ) as encode_images: with mock.patch("ComfyUI_LLM_API.node.create_openai_client") as create_client: with mock.patch("ComfyUI_LLM_API.node.logger.warning"): response = _run_node( provider_profile="custom_text_only", images=object(), ) self.assertEqual(response[0], "") self.assertIn("does not support images", response[1]) encode_images.assert_not_called() create_client.assert_not_called() def test_blank_openai_base_url_uses_profile_default_and_override_wins(self): with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()) as create_client: with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", return_value=_completion("ok"), ): _run_node(provider_profile="openai", api_baseurl="") _run_node( provider_profile="openai", api_baseurl="https://override.example/v1", ) self.assertEqual( create_client.call_args_list[0].kwargs["base_url"], "https://api.openai.com/v1", ) self.assertEqual( create_client.call_args_list[1].kwargs["base_url"], "https://override.example/v1", ) def test_image_success_encodes_once_through_retry_message_factory(self): def fake_create_chat_completion(**kwargs): factory = kwargs["image_retry_message_factory"] messages = factory(kwargs["image_retry_limits"][0]) self.assertTrue(kwargs["has_images"]) self.assertEqual(messages[-1]["content"][1]["image_url"]["url"], "data:image/png;base64,abc") return _completion("vision hello") with mock.patch( "ComfyUI_LLM_API.node.encode_comfy_image_batch_to_data_urls", return_value=["data:image/png;base64,abc"], ) as encode_images: with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", side_effect=fake_create_chat_completion, ): response = _run_node(images=object()) self.assertEqual(response, ("vision hello", "")) encode_images.assert_called_once() self.assertEqual(encode_images.call_args.kwargs["max_bytes"], 1_500_000) self.assertEqual(encode_images.call_args.kwargs["max_side"], 1280) def test_custom_image_preset_forwards_custom_mb_and_side(self): def fake_create_chat_completion(**kwargs): kwargs["image_retry_message_factory"](kwargs["image_retry_limits"][0]) return _completion("vision hello") with mock.patch( "ComfyUI_LLM_API.node.encode_comfy_image_batch_to_data_urls", return_value=["data:image/png;base64,abc"], ) as encode_images: with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", side_effect=fake_create_chat_completion, ): response = _run_node( images=object(), image_preset="custom", custom_image_max_mb=2.5, custom_image_max_side=1400, ) self.assertEqual(response, ("vision hello", "")) self.assertEqual(encode_images.call_args.kwargs["max_bytes"], 2_500_000) self.assertEqual(encode_images.call_args.kwargs["max_side"], 1400) def test_value_error_returns_sanitized_error_string(self): with mock.patch("ComfyUI_LLM_API.node.logger.warning") as warning: with mock.patch( "ComfyUI_LLM_API.node.create_openai_client", side_effect=ValueError("bad sk-secret"), ): response = _run_node() warning.assert_called_once() self.assertNotIn("sk-secret", " ".join(str(arg) for arg in warning.call_args.args)) self.assertEqual(response[0], "") self.assertTrue(response[1].startswith("LLM API Error: ")) self.assertNotIn("sk-secret", response[1]) def test_empty_completion_returns_error_output(self): with mock.patch("ComfyUI_LLM_API.node.create_openai_client", return_value=object()): with mock.patch( "ComfyUI_LLM_API.node.create_chat_completion", return_value=SimpleNamespace(choices=[]), ): with mock.patch("ComfyUI_LLM_API.node.logger.warning"): response = _run_node() self.assertEqual(response[0], "") self.assertTrue(response[1].startswith("LLM API Error: ")) self.assertIn("no choices", response[1].lower()) if __name__ == "__main__": unittest.main()