109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
import io
|
|
import unittest
|
|
from contextlib import redirect_stdout
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
from PIL import Image
|
|
|
|
try:
|
|
from . import test_path
|
|
except ImportError:
|
|
import test_path
|
|
|
|
|
|
def _completion(text):
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content=text))]
|
|
)
|
|
|
|
|
|
class SmokeProviderScriptTests(unittest.TestCase):
|
|
def test_missing_env_key_exits_with_readable_message(self):
|
|
from ComfyUI_LLM_API.scripts import smoke_provider
|
|
|
|
output = io.StringIO()
|
|
with mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LLM_API_PROVIDER_PROFILE": "openai",
|
|
"LLM_API_MODEL": "demo-model",
|
|
},
|
|
clear=True,
|
|
):
|
|
with redirect_stdout(output):
|
|
exit_code = smoke_provider.main()
|
|
|
|
self.assertEqual(exit_code, 1)
|
|
self.assertIn("LLM_API_KEY", output.getvalue())
|
|
|
|
def test_text_smoke_passes_fake_client_without_printing_key(self):
|
|
from ComfyUI_LLM_API.scripts import smoke_provider
|
|
|
|
output = io.StringIO()
|
|
with mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LLM_API_PROVIDER_PROFILE": "openai",
|
|
"LLM_API_KEY": "sk-secret",
|
|
"LLM_API_MODEL": "demo-model",
|
|
},
|
|
clear=True,
|
|
):
|
|
with mock.patch(
|
|
"ComfyUI_LLM_API.scripts.smoke_provider.create_openai_client",
|
|
return_value=object(),
|
|
) as create_client:
|
|
with mock.patch(
|
|
"ComfyUI_LLM_API.scripts.smoke_provider.create_chat_completion",
|
|
return_value=_completion("hello"),
|
|
) as create_completion:
|
|
with redirect_stdout(output):
|
|
exit_code = smoke_provider.main()
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
self.assertEqual(create_client.call_args.kwargs["api_key"], "sk-secret")
|
|
self.assertFalse(create_completion.call_args.kwargs["has_images"])
|
|
self.assertIn("text: ok", output.getvalue())
|
|
self.assertIn("image: skipped", output.getvalue())
|
|
self.assertNotIn("sk-secret", output.getvalue())
|
|
|
|
def test_image_smoke_uses_image_path_only_when_set(self):
|
|
from ComfyUI_LLM_API.scripts import smoke_provider
|
|
|
|
image = Image.new("RGB", (8, 8), (20, 40, 60))
|
|
output = io.StringIO()
|
|
|
|
with mock.patch.dict(
|
|
"os.environ",
|
|
{
|
|
"LLM_API_PROVIDER_PROFILE": "openai",
|
|
"LLM_API_KEY": "sk-secret",
|
|
"LLM_API_MODEL": "demo-model",
|
|
"LLM_API_IMAGE_PATH": "demo.png",
|
|
},
|
|
clear=True,
|
|
):
|
|
with mock.patch(
|
|
"ComfyUI_LLM_API.scripts.smoke_provider.create_openai_client",
|
|
return_value=object(),
|
|
):
|
|
with mock.patch(
|
|
"ComfyUI_LLM_API.scripts.smoke_provider.create_chat_completion",
|
|
return_value=_completion("hello"),
|
|
) as create_completion:
|
|
with mock.patch(
|
|
"ComfyUI_LLM_API.scripts.smoke_provider.Image.open",
|
|
return_value=image,
|
|
):
|
|
with redirect_stdout(output):
|
|
exit_code = smoke_provider.main()
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
self.assertTrue(create_completion.call_args_list[-1].kwargs["has_images"])
|
|
self.assertIn("image: ok", output.getvalue())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|