From 223ec7e15b57ca48fcb93ae7bdceb479d2b59d2a Mon Sep 17 00:00:00 2001
From: by-ae <25704492+by-ae@users.noreply.github.com>
Date: Fri, 23 Oct 2026 21:15:26 +0000
Subject: [PATCH] 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>
---
 .github/workflows/ci.yml |   4 +-
 tests/conftest.py        |  75 ++++++++++++
 tests/test_nodes.py      | 242 +++++++++++++++++++++++++++++++++++++++
 tests/test_routes.py     | 192 +++++++++++++++++++++++++++++++
 4 files changed, 511 insertions(+), 2 deletions(-)
 create mode 100644 tests/conftest.py
 create mode 100644 tests/test_nodes.py
 create mode 100644 tests/test_routes.py

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f0fc02e..32418cc 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -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
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..539b50c
--- /dev/null
+++ b/tests/conftest.py
@@ -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
diff --git a/tests/test_nodes.py b/tests/test_nodes.py
new file mode 100644
index 0000000..37ab9e6
--- /dev/null
+++ b/tests/test_nodes.py
@@ -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()
diff --git a/tests/test_routes.py b/tests/test_routes.py
new file mode 100644
index 0000000..5e09f58
--- /dev/null
+++ b/tests/test_routes.py
@@ -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()