Add comprehensive test coverage for routes and nodes
Tests cover: - routes.py: Input validation, path traversal protection, JSON error handling, GUID sanitization (12 new tests) - nodes/ideogram4_nodes.py: Prompt assembly, bbox normalization, palette handling, JSON serialization, input coercion (28 new tests) All tests use mocks via conftest.py to avoid ComfyUI dependencies. Tests import the repo as a proper package so routes.py relative imports work unchanged. CI: Install aiohttp for route tests. Co-authored-by: Cursor Agent <cursoragent@cursor.com>
This commit is contained in:
by-ae
@@ -24,7 +24,7 @@ jobs:
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install flake8
|
||||
pip install flake8 aiohttp pytest pytest-asyncio
|
||||
|
||||
- name: Lint with flake8
|
||||
run: |
|
||||
@@ -43,7 +43,7 @@ jobs:
|
||||
|
||||
- name: Run unit tests
|
||||
run: |
|
||||
python -m unittest discover -s tests -p 'test_*.py' -v
|
||||
python -m pytest tests/ -v
|
||||
|
||||
javascript-check:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
"""pytest configuration for ComfyUI-ByAE-IdeogramTools tests."""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock ComfyUI dependencies before any imports
|
||||
sys.modules['folder_paths'] = MagicMock()
|
||||
sys.modules['torch'] = MagicMock()
|
||||
sys.modules['numpy'] = MagicMock()
|
||||
sys.modules['PIL'] = MagicMock()
|
||||
sys.modules['PIL.Image'] = MagicMock()
|
||||
sys.modules['PIL.ImageDraw'] = MagicMock()
|
||||
sys.modules['PIL.ImageFont'] = MagicMock()
|
||||
sys.modules['comfy_api'] = MagicMock()
|
||||
sys.modules['comfy_api.latest'] = MagicMock()
|
||||
sys.modules['comfy_api.latest.io'] = MagicMock()
|
||||
sys.modules['server'] = MagicMock()
|
||||
|
||||
# Make the repo root importable as a package
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
parent_dir = os.path.dirname(repo_root)
|
||||
package_name = os.path.basename(repo_root).replace('-', '_')
|
||||
|
||||
if parent_dir not in sys.path:
|
||||
sys.path.insert(0, parent_dir)
|
||||
|
||||
# Create a pseudo-package that wraps the repo
|
||||
import importlib.util
|
||||
import types
|
||||
|
||||
# Create package module
|
||||
pkg = types.ModuleType(package_name)
|
||||
pkg.__path__ = [repo_root]
|
||||
pkg.__file__ = os.path.join(repo_root, '__init__.py')
|
||||
sys.modules[package_name] = pkg
|
||||
|
||||
# Import and register store module
|
||||
store_spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.store",
|
||||
os.path.join(repo_root, 'store.py')
|
||||
)
|
||||
store_module = importlib.util.module_from_spec(store_spec)
|
||||
sys.modules[f"{package_name}.store"] = store_module
|
||||
store_spec.loader.exec_module(store_module)
|
||||
|
||||
# Import and register nodes package
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = [os.path.join(repo_root, 'nodes')]
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
# Import and register ideogram4_nodes
|
||||
nodes_spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.nodes.ideogram4_nodes",
|
||||
os.path.join(repo_root, 'nodes', 'ideogram4_nodes.py')
|
||||
)
|
||||
nodes_module = importlib.util.module_from_spec(nodes_spec)
|
||||
sys.modules[f"{package_name}.nodes.ideogram4_nodes"] = nodes_module
|
||||
nodes_spec.loader.exec_module(nodes_module)
|
||||
|
||||
# Import and register routes with the package context
|
||||
routes_spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.routes",
|
||||
os.path.join(repo_root, 'routes.py'),
|
||||
submodule_search_locations=[repo_root]
|
||||
)
|
||||
routes_module = importlib.util.module_from_spec(routes_spec)
|
||||
sys.modules[f"{package_name}.routes"] = routes_module
|
||||
routes_spec.loader.exec_module(routes_module)
|
||||
|
||||
# Make them available without package prefix in test files
|
||||
sys.modules['store'] = store_module
|
||||
sys.modules['routes'] = routes_module
|
||||
sys.modules['nodes'] = nodes_pkg
|
||||
sys.modules['nodes.ideogram4_nodes'] = nodes_module
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Unit tests for nodes/ideogram4_nodes.py prompt assembly and validation."""
|
||||
|
||||
import json
|
||||
import unittest
|
||||
|
||||
# conftest.py handles mocking and path setup
|
||||
from nodes.ideogram4_nodes import (
|
||||
sanitize_bg_guid, _norm_bbox, _palette, _dumps, _parse_json_list,
|
||||
_caption_to_boxes, _coerce_layout_int, _coerce_opacity
|
||||
)
|
||||
|
||||
|
||||
class TestPromptHelpers(unittest.TestCase):
|
||||
def test_sanitize_bg_guid_valid(self):
|
||||
valid_guid = "12345678-1234-4abc-8abc-123456789012"
|
||||
self.assertEqual(sanitize_bg_guid(valid_guid), valid_guid)
|
||||
|
||||
# Test uppercase
|
||||
upper_guid = "12345678-1234-4ABC-8ABC-123456789012"
|
||||
self.assertEqual(sanitize_bg_guid(upper_guid), upper_guid)
|
||||
|
||||
def test_sanitize_bg_guid_invalid(self):
|
||||
with self.assertRaises(ValueError):
|
||||
sanitize_bg_guid("not-a-guid")
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
sanitize_bg_guid("")
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
sanitize_bg_guid("12345678-1234-4abc-8abc") # Too short
|
||||
|
||||
# Path traversal attempt
|
||||
with self.assertRaises(ValueError):
|
||||
sanitize_bg_guid("../../etc/passwd")
|
||||
|
||||
def test_norm_bbox_basic(self):
|
||||
box = {"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4}
|
||||
result = _norm_bbox(box)
|
||||
self.assertEqual(result, [200, 100, 600, 400])
|
||||
|
||||
def test_norm_bbox_swap_coords_if_inverted(self):
|
||||
# Test ymin > ymax case
|
||||
box = {"x": 0.1, "y": 0.6, "w": 0.3, "h": -0.4}
|
||||
result = _norm_bbox(box)
|
||||
# Should swap ymin and ymax
|
||||
self.assertEqual(result[0], min(result[0], result[2]))
|
||||
self.assertEqual(result[2], max(result[0], result[2]))
|
||||
|
||||
def test_norm_bbox_clamps_to_1000(self):
|
||||
box = {"x": 0.0, "y": 0.0, "w": 2.0, "h": 2.0}
|
||||
result = _norm_bbox(box)
|
||||
self.assertTrue(all(0 <= v <= 1000 for v in result))
|
||||
|
||||
def test_norm_bbox_missing_fields(self):
|
||||
# Should handle missing fields gracefully
|
||||
box = {}
|
||||
result = _norm_bbox(box)
|
||||
self.assertEqual(result, [0, 0, 0, 0])
|
||||
|
||||
def test_palette_filters_empty_strings(self):
|
||||
colors = {"c1": "#FF0000", "c2": "", "c3": "#00FF00", "c4": None}
|
||||
result = _palette(colors)
|
||||
self.assertEqual(result, ["#FF0000", "#00FF00"])
|
||||
|
||||
def test_palette_uppercase(self):
|
||||
colors = ["#ff0000", "#00ff00"]
|
||||
result = _palette(colors)
|
||||
self.assertEqual(result, ["#FF0000", "#00FF00"])
|
||||
|
||||
def test_dumps_string(self):
|
||||
result = _dumps("hello")
|
||||
self.assertEqual(result, '"hello"')
|
||||
|
||||
def test_dumps_list_simple(self):
|
||||
result = _dumps([1, 2, 3])
|
||||
self.assertEqual(result, "[1, 2, 3]")
|
||||
|
||||
def test_dumps_list_nested(self):
|
||||
result = _dumps([{"a": 1}, {"b": 2}])
|
||||
self.assertIn("[\n", result)
|
||||
self.assertIn('"a": 1', result)
|
||||
|
||||
def test_dumps_dict(self):
|
||||
result = _dumps({"key": "value", "num": 123})
|
||||
self.assertIn('"key": "value"', result)
|
||||
self.assertIn('"num": 123', result)
|
||||
|
||||
def test_dumps_unicode(self):
|
||||
result = _dumps({"text": "Hello 世界"})
|
||||
self.assertIn("世界", result)
|
||||
|
||||
def test_parse_json_list_valid(self):
|
||||
result = _parse_json_list('[{"x": 1}, {"y": 2}]')
|
||||
self.assertEqual(result, [{"x": 1}, {"y": 2}])
|
||||
|
||||
def test_parse_json_list_invalid(self):
|
||||
result = _parse_json_list("not json")
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_parse_json_list_not_list(self):
|
||||
result = _parse_json_list('{"key": "value"}')
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_parse_json_list_empty(self):
|
||||
result = _parse_json_list("")
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_caption_to_boxes_valid(self):
|
||||
caption = {
|
||||
"compositional_deconstruction": {
|
||||
"elements": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Hello",
|
||||
"desc": "A greeting",
|
||||
"bbox": [100, 200, 300, 400],
|
||||
"color_palette": ["#FF0000"]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
result = _caption_to_boxes(caption)
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertEqual(result[0]["type"], "text")
|
||||
self.assertEqual(result[0]["text"], "Hello")
|
||||
self.assertAlmostEqual(result[0]["x"], 0.2)
|
||||
self.assertAlmostEqual(result[0]["y"], 0.1)
|
||||
|
||||
def test_caption_to_boxes_missing_bbox(self):
|
||||
caption = {
|
||||
"compositional_deconstruction": {
|
||||
"elements": [
|
||||
{"type": "obj", "desc": "No bbox"}
|
||||
]
|
||||
}
|
||||
}
|
||||
result = _caption_to_boxes(caption)
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertTrue(result[0]["nobbox"])
|
||||
|
||||
def test_caption_to_boxes_invalid_bbox(self):
|
||||
caption = {
|
||||
"compositional_deconstruction": {
|
||||
"elements": [
|
||||
{"type": "obj", "desc": "Bad bbox", "bbox": "not a list"}
|
||||
]
|
||||
}
|
||||
}
|
||||
result = _caption_to_boxes(caption)
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertTrue(result[0]["nobbox"])
|
||||
|
||||
def test_caption_to_boxes_empty(self):
|
||||
result = _caption_to_boxes({})
|
||||
self.assertEqual(result, [])
|
||||
|
||||
def test_coerce_layout_int_valid(self):
|
||||
self.assertEqual(_coerce_layout_int(512), 512)
|
||||
self.assertEqual(_coerce_layout_int("1024"), 1024)
|
||||
self.assertEqual(_coerce_layout_int(2048.7), 2048)
|
||||
|
||||
def test_coerce_layout_int_invalid(self):
|
||||
self.assertEqual(_coerce_layout_int(0), 1024)
|
||||
self.assertEqual(_coerce_layout_int(-100), 1024)
|
||||
self.assertEqual(_coerce_layout_int("not a number"), 1024)
|
||||
self.assertEqual(_coerce_layout_int(None), 1024)
|
||||
|
||||
def test_coerce_opacity_valid(self):
|
||||
self.assertEqual(_coerce_opacity(0.5), 0.5)
|
||||
self.assertEqual(_coerce_opacity(0), 0.0)
|
||||
self.assertEqual(_coerce_opacity(1), 1.0)
|
||||
self.assertEqual(_coerce_opacity("0.75"), 0.75)
|
||||
|
||||
def test_coerce_opacity_clamping(self):
|
||||
self.assertEqual(_coerce_opacity(-1), 0.0)
|
||||
self.assertEqual(_coerce_opacity(2), 1.0)
|
||||
self.assertEqual(_coerce_opacity(1.5), 1.0)
|
||||
|
||||
def test_coerce_opacity_invalid(self):
|
||||
self.assertEqual(_coerce_opacity("not a number"), 0.5)
|
||||
self.assertEqual(_coerce_opacity(None), 0.5)
|
||||
self.assertEqual(_coerce_opacity(float('nan')), 0.5)
|
||||
|
||||
|
||||
class TestPromptAssembly(unittest.TestCase):
|
||||
"""Test full prompt assembly logic without ComfyUI runtime."""
|
||||
|
||||
def test_build_minimal_caption(self):
|
||||
# Simulate minimal inputs
|
||||
boxes = []
|
||||
background = "A simple scene"
|
||||
|
||||
caption = {
|
||||
"compositional_deconstruction": {
|
||||
"background": background,
|
||||
"elements": []
|
||||
}
|
||||
}
|
||||
|
||||
json_str = _dumps(caption)
|
||||
self.assertIn('"background":', json_str)
|
||||
self.assertIn('"A simple scene"', json_str)
|
||||
|
||||
def test_build_caption_with_elements(self):
|
||||
boxes = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "HELLO",
|
||||
"desc": "Large text",
|
||||
"x": 0.1, "y": 0.2, "w": 0.3, "h": 0.4,
|
||||
"palette": ["#FF0000", "#00FF00"]
|
||||
},
|
||||
{
|
||||
"type": "obj",
|
||||
"desc": "A circle",
|
||||
"x": 0.5, "y": 0.5, "w": 0.2, "h": 0.2,
|
||||
"palette": []
|
||||
}
|
||||
]
|
||||
|
||||
elements = []
|
||||
for box in boxes:
|
||||
elem = {"type": box["type"]}
|
||||
if "nobbox" not in box:
|
||||
elem["bbox"] = _norm_bbox(box)
|
||||
if box["type"] == "text":
|
||||
elem["text"] = box["text"]
|
||||
elem["desc"] = box["desc"]
|
||||
pal = _palette(box["palette"])
|
||||
if pal:
|
||||
elem["color_palette"] = pal
|
||||
elements.append(elem)
|
||||
|
||||
self.assertEqual(len(elements), 2)
|
||||
self.assertEqual(elements[0]["type"], "text")
|
||||
self.assertEqual(elements[0]["text"], "HELLO")
|
||||
self.assertEqual(elements[1]["type"], "obj")
|
||||
self.assertEqual(len(elements[0]["color_palette"]), 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""Unit tests for the routes module."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import AioHTTPTestCase
|
||||
|
||||
# conftest.py handles mocking and path setup
|
||||
import store
|
||||
import routes
|
||||
from nodes.ideogram4_nodes import sanitize_bg_guid, bg_temp_filename, bg_temp_path
|
||||
|
||||
|
||||
class TestRoutes(AioHTTPTestCase):
|
||||
async def get_application(self):
|
||||
app = web.Application()
|
||||
return app
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
# Get the actual mocked folder_paths from sys.modules (created by conftest)
|
||||
import sys
|
||||
self.mock_folder_paths = sys.modules['folder_paths']
|
||||
self.mock_folder_paths.get_user_directory.return_value = self.temp_dir
|
||||
self.mock_folder_paths.get_temp_directory.return_value = self.temp_dir
|
||||
|
||||
def tearDown(self):
|
||||
super().tearDown()
|
||||
import shutil
|
||||
if os.path.exists(self.temp_dir):
|
||||
shutil.rmtree(self.temp_dir)
|
||||
|
||||
async def test_get_stores_success(self):
|
||||
# Pre-populate store
|
||||
store.save_entry("characters", "TestChar", {"desc": "A test"})
|
||||
|
||||
request = MagicMock(spec=web.Request)
|
||||
response = await routes.get_stores(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("characters", data)
|
||||
self.assertIn("TestChar", data["characters"])
|
||||
|
||||
async def test_get_stores_empty(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
response = await routes.get_stores(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
data = json.loads(response.body)
|
||||
self.assertEqual(data["characters"], {})
|
||||
self.assertEqual(data["objects"], {})
|
||||
self.assertEqual(data["scenes"], {})
|
||||
|
||||
async def test_put_store_entry_success(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "characters", "name": "Hero"}
|
||||
request.json = AsyncMock(return_value={"desc": "Main character", "palette": ["#FF0000"]})
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
data = json.loads(response.body)
|
||||
self.assertTrue(data["ok"])
|
||||
|
||||
# Verify it was saved
|
||||
bucket = store.load_bucket("characters")
|
||||
self.assertIn("Hero", bucket)
|
||||
self.assertEqual(bucket["Hero"]["desc"], "Main character")
|
||||
|
||||
async def test_put_store_entry_invalid_bucket(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "invalid_bucket", "name": "Test"}
|
||||
request.json = AsyncMock(return_value={})
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("error", data)
|
||||
self.assertIn("Invalid bucket", data["error"])
|
||||
|
||||
async def test_put_store_entry_invalid_name(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "characters", "name": ""}
|
||||
request.json = AsyncMock(return_value={})
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("error", data)
|
||||
|
||||
async def test_put_store_entry_name_too_long(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "characters", "name": "a" * 200}
|
||||
request.json = AsyncMock(return_value={})
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("error", data)
|
||||
|
||||
async def test_put_store_entry_invalid_json(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "characters", "name": "Test"}
|
||||
request.json = AsyncMock(side_effect=json.JSONDecodeError("", "", 0))
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("error", data)
|
||||
self.assertIn("Invalid JSON", data["error"])
|
||||
|
||||
async def test_put_store_entry_non_object_body(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"bucket": "characters", "name": "Test"}
|
||||
request.json = AsyncMock(return_value=["array", "not", "object"])
|
||||
|
||||
response = await routes.put_store_entry(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertIn("error", data)
|
||||
self.assertIn("JSON object", data["error"])
|
||||
|
||||
async def test_get_bg_stat_valid_guid(self):
|
||||
# Create a background file
|
||||
guid = "12345678-1234-4abc-8abc-123456789012"
|
||||
bg_path = os.path.join(self.temp_dir, f"byae_ideogram_bg_{guid}.png")
|
||||
with open(bg_path, "wb") as f:
|
||||
f.write(b"fake png data")
|
||||
|
||||
# Mock the bg_temp_path to return our test path
|
||||
with patch('routes.bg_temp_path', return_value=bg_path):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"guid": guid}
|
||||
|
||||
response = await routes.get_bg_stat(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
data = json.loads(response.body)
|
||||
self.assertTrue(data["exists"])
|
||||
self.assertGreater(data["mtime"], 0)
|
||||
self.assertIn("url", data)
|
||||
|
||||
async def test_get_bg_stat_file_not_exists(self):
|
||||
guid = "12345678-1234-4abc-8abc-123456789012"
|
||||
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"guid": guid}
|
||||
|
||||
response = await routes.get_bg_stat(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
data = json.loads(response.body)
|
||||
self.assertFalse(data["exists"])
|
||||
self.assertEqual(data["mtime"], 0)
|
||||
|
||||
async def test_get_bg_stat_invalid_guid(self):
|
||||
# Test path traversal attempt
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"guid": "../../../etc/passwd"}
|
||||
|
||||
response = await routes.get_bg_stat(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertFalse(data["exists"])
|
||||
self.assertIn("error", data)
|
||||
|
||||
async def test_get_bg_stat_malformed_guid(self):
|
||||
request = MagicMock(spec=web.Request)
|
||||
request.match_info = {"guid": "not-a-valid-guid"}
|
||||
|
||||
response = await routes.get_bg_stat(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
data = json.loads(response.body)
|
||||
self.assertFalse(data["exists"])
|
||||
self.assertIn("error", data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user