Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
116 lines
4.6 KiB
Python
116 lines
4.6 KiB
Python
from __future__ import annotations
|
||
|
||
import importlib.util
|
||
import math
|
||
from pathlib import Path
|
||
import sys
|
||
import types
|
||
import unittest
|
||
from unittest import mock
|
||
|
||
|
||
PLUGIN_ROOT = Path(__file__).resolve().parents[1]
|
||
|
||
|
||
def _load_source_module(name: str, path: Path):
|
||
spec = importlib.util.spec_from_file_location(name, path)
|
||
if spec is None or spec.loader is None:
|
||
raise RuntimeError(f"无法加载测试模块：{path}")
|
||
module = importlib.util.module_from_spec(spec)
|
||
sys.modules[name] = module
|
||
spec.loader.exec_module(module)
|
||
return module
|
||
|
||
|
||
def _load_prompt_module():
|
||
root_package = types.ModuleType("comfyui_o1key")
|
||
root_package.__path__ = [str(PLUGIN_ROOT)]
|
||
sys.modules["comfyui_o1key"] = root_package
|
||
|
||
for package_name, directory in (
|
||
("comfyui_o1key.utils", PLUGIN_ROOT / "utils"),
|
||
("comfyui_o1key.nodes", PLUGIN_ROOT / "nodes"),
|
||
):
|
||
package = types.ModuleType(package_name)
|
||
package.__path__ = [str(directory)]
|
||
sys.modules[package_name] = package
|
||
|
||
_load_source_module(
|
||
"comfyui_o1key.utils.image_utils",
|
||
PLUGIN_ROOT / "utils" / "image_utils.py",
|
||
)
|
||
return _load_source_module(
|
||
"comfyui_o1key.nodes.prompt_multi_function",
|
||
PLUGIN_ROOT / "nodes" / "prompt_multi_function.py",
|
||
)
|
||
|
||
|
||
prompt_module = _load_prompt_module()
|
||
O1keyPromptMultiFunction = prompt_module.O1keyPromptMultiFunction
|
||
|
||
|
||
def _prompts(count: int) -> str:
|
||
return "\n---\n".join(f"提示词 {index}" for index in range(1, count + 1))
|
||
|
||
|
||
class PromptMultiFunctionTests(unittest.TestCase):
|
||
def setUp(self):
|
||
self.node = O1keyPromptMultiFunction()
|
||
|
||
def test_schema_appends_multi_select_widgets_without_changing_legacy_positions(self):
|
||
required = O1keyPromptMultiFunction.INPUT_TYPES()["required"]
|
||
self.assertEqual(list(required), ["提示词", "功能", "抽取数量", "指定序号"])
|
||
self.assertEqual(required["功能"][0], ["全部使用", "随机抽取n套", "指定序号"])
|
||
self.assertEqual(required["抽取数量"][1]["default"], 3)
|
||
self.assertEqual(required["指定序号"][1]["default"], "1,2,3")
|
||
|
||
def test_legacy_random_one_mode_still_accepts_the_old_two_argument_call(self):
|
||
with mock.patch.object(prompt_module.random, "choice", return_value="提示词 2"):
|
||
result = self.node.process(_prompts(3), "随机抽取1套")
|
||
self.assertEqual(result, ("提示词 2",))
|
||
|
||
def test_random_many_samples_without_replacement_and_restores_source_order(self):
|
||
with mock.patch.object(prompt_module.random, "sample", return_value=[4, 1, 3]) as sample:
|
||
result = self.node.process(_prompts(5), "随机抽取n套", 3)
|
||
sample.assert_called_once()
|
||
population, count = sample.call_args.args
|
||
self.assertEqual(list(population), [0, 1, 2, 3, 4])
|
||
self.assertEqual(count, 3)
|
||
self.assertEqual(result, ("提示词 2\n---\n提示词 4\n---\n提示词 5",))
|
||
|
||
def test_random_many_rejects_a_count_larger_than_the_prompt_set(self):
|
||
with self.assertRaisesRegex(ValueError, "抽取数量 4 超过当前提示词总数 3"):
|
||
self.node.process(_prompts(3), "随机抽取n套", 4)
|
||
|
||
def test_selected_indices_support_common_separators_ranges_and_input_order(self):
|
||
result = self.node.process(
|
||
_prompts(5),
|
||
"指定序号",
|
||
指定序号="5，2 3-4",
|
||
)
|
||
self.assertEqual(
|
||
result,
|
||
("提示词 5\n---\n提示词 2\n---\n提示词 3\n---\n提示词 4",),
|
||
)
|
||
|
||
def test_selected_indices_reject_duplicates_and_out_of_range_values(self):
|
||
with self.assertRaisesRegex(ValueError, "序号 2 重复"):
|
||
self.node.process(_prompts(3), "指定序号", 指定序号="2,2")
|
||
with self.assertRaisesRegex(ValueError, "序号 4 超出范围"):
|
||
self.node.process(_prompts(3), "指定序号", 指定序号="1,4")
|
||
|
||
def test_only_random_modes_disable_execution_caching(self):
|
||
self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取1套")))
|
||
self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取多套")))
|
||
self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取n套")))
|
||
selected_key = O1keyPromptMultiFunction.IS_CHANGED("a", "指定序号", 3, "1,3")
|
||
self.assertEqual(selected_key, "指定序号|1,3|a")
|
||
self.assertEqual(
|
||
O1keyPromptMultiFunction.IS_CHANGED("a", "全部使用", 9, "2"),
|
||
"全部使用|a",
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|