176 lines
7.1 KiB
Python
176 lines
7.1 KiB
Python
"""
|
|
Tests for R8: Smart Workflow Truncation
|
|
"""
|
|
import unittest
|
|
import sys
|
|
import os
|
|
import json
|
|
|
|
# --- PATH SETUP ---
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.abspath(os.path.join(current_dir, ".."))
|
|
if project_root not in sys.path:
|
|
sys.path.insert(0, project_root)
|
|
|
|
from truncate_workflow import truncate_workflow_smart
|
|
|
|
|
|
class TestWorkflowTruncation(unittest.TestCase):
|
|
"""Tests for R8 smart workflow truncation."""
|
|
|
|
def test_small_workflow_unchanged(self):
|
|
"""Small workflows should pass through unchanged."""
|
|
small_workflow = json.dumps({"nodes": [{"id": 1, "type": "KSampler"}]})
|
|
result, meta = truncate_workflow_smart(small_workflow, max_chars=1000)
|
|
|
|
self.assertEqual(meta["truncation_method"], "none")
|
|
self.assertEqual(result, small_workflow)
|
|
|
|
def test_empty_workflow(self):
|
|
"""Empty workflow should return empty string."""
|
|
result, meta = truncate_workflow_smart("", max_chars=100)
|
|
|
|
self.assertEqual(result, "")
|
|
self.assertEqual(meta["truncation_method"], "none")
|
|
|
|
def test_large_workflow_truncated(self):
|
|
"""Large workflows should be truncated."""
|
|
# Create a large workflow with many nodes
|
|
nodes = [{"id": i, "type": f"Node{i}", "properties": {"data": "x" * 100}} for i in range(50)]
|
|
large_workflow = json.dumps({"nodes": nodes})
|
|
|
|
result, meta = truncate_workflow_smart(large_workflow, max_chars=500)
|
|
|
|
self.assertLess(len(result), len(large_workflow))
|
|
self.assertIn(meta["truncation_method"], ["property_pruning", "node_removal", "char_slice"])
|
|
self.assertIn("original_nodes", meta)
|
|
|
|
def test_returns_valid_json_or_readable_string(self):
|
|
"""Output should be valid JSON or at least readable."""
|
|
nodes = [{"id": i, "type": f"Node{i}"} for i in range(100)]
|
|
workflow = json.dumps({"nodes": nodes})
|
|
|
|
result, meta = truncate_workflow_smart(workflow, max_chars=200)
|
|
|
|
# Should either be valid JSON or readable text
|
|
try:
|
|
json.loads(result)
|
|
valid_json = True
|
|
except json.JSONDecodeError:
|
|
valid_json = False
|
|
# If not valid JSON, should at least have truncation marker
|
|
self.assertIn("truncated", result.lower())
|
|
|
|
def test_preserves_error_node(self):
|
|
"""Error node should be preserved when specified."""
|
|
nodes = [
|
|
{"id": "1", "type": "NodeA", "inputs": {}},
|
|
{"id": "2", "type": "NodeB", "inputs": {"a": ["1", 0]}},
|
|
{"id": "3", "type": "NodeC", "inputs": {"b": ["2", 0]}},
|
|
{"id": "4", "type": "NodeD", "inputs": {}},
|
|
{"id": "5", "type": "NodeE", "inputs": {}},
|
|
]
|
|
# Add some bulk to force truncation
|
|
for node in nodes:
|
|
node["properties"] = {"large_data": "x" * 500}
|
|
|
|
workflow = json.dumps({"nodes": nodes})
|
|
|
|
# Truncate with error node = "3"
|
|
result, meta = truncate_workflow_smart(workflow, error_node_id="3", max_chars=800)
|
|
|
|
# Check that error node 3 is preserved (if node_removal was used)
|
|
if meta["truncation_method"] == "node_removal":
|
|
self.assertTrue(meta.get("error_node_preserved", True))
|
|
|
|
def test_invalid_json_fallback(self):
|
|
"""Invalid JSON should fall back to character truncation."""
|
|
invalid_workflow = "This is not JSON at all, just plain text" * 100
|
|
|
|
result, meta = truncate_workflow_smart(invalid_workflow, max_chars=100)
|
|
|
|
self.assertEqual(meta["truncation_method"], "char_slice")
|
|
self.assertLessEqual(len(result), 100)
|
|
|
|
def test_metadata_contains_required_fields(self):
|
|
"""Metadata should contain required fields."""
|
|
workflow = json.dumps({"nodes": [{"id": 1}]})
|
|
_, meta = truncate_workflow_smart(workflow)
|
|
|
|
self.assertIn("truncation_method", meta)
|
|
self.assertIn("original_length", meta)
|
|
|
|
def test_property_pruning_preserves_node_count(self):
|
|
"""Property pruning should keep all nodes."""
|
|
nodes = [
|
|
{"id": 1, "type": "A", "widgets_values": [1, 2, 3], "pos": [0, 0], "size": [100, 100]},
|
|
{"id": 2, "type": "B", "widgets_values": [4, 5, 6], "pos": [100, 0], "size": [100, 100]},
|
|
]
|
|
workflow = json.dumps({"nodes": nodes})
|
|
|
|
result, meta = truncate_workflow_smart(workflow, max_chars=200)
|
|
|
|
if meta["truncation_method"] == "property_pruning":
|
|
self.assertEqual(meta["original_nodes"], meta["kept_nodes"])
|
|
|
|
def test_execution_id_error_node_matches_local_node_id(self):
|
|
"""Execution IDs from subgraphs should still preserve the target node."""
|
|
nodes = []
|
|
for i in range(1, 8):
|
|
nodes.append({
|
|
"id": str(i),
|
|
"type": f"Node{i}",
|
|
"inputs": {"x": [str(i - 1), 0]} if i > 1 else {},
|
|
"properties": {"large_data": "x" * 400},
|
|
})
|
|
workflow = json.dumps({"nodes": nodes, "links": []})
|
|
|
|
result, meta = truncate_workflow_smart(
|
|
workflow,
|
|
error_node_id="65:70:6",
|
|
max_chars=700,
|
|
)
|
|
|
|
self.assertEqual(meta["truncation_method"], "node_removal")
|
|
self.assertTrue(meta.get("error_node_preserved"))
|
|
parsed = json.loads(result)
|
|
self.assertTrue(any(str(node.get("id")) == "6" for node in parsed["nodes"]))
|
|
|
|
def test_ui_workflow_pruning_preserves_metadata_and_filters_links(self):
|
|
"""UI-format workflow pruning should preserve top-level metadata and valid links."""
|
|
nodes = [
|
|
{"id": "1", "type": "Loader", "inputs": {}, "properties": {"blob": "x" * 400}},
|
|
{"id": "2", "type": "KSampler", "inputs": {"model": ["1", 0]}, "properties": {"blob": "y" * 400}},
|
|
{"id": "3", "type": "SaveImage", "inputs": {"images": ["2", 0]}, "properties": {"blob": "z" * 400}},
|
|
{"id": "4", "type": "Unused", "inputs": {}, "properties": {"blob": "w" * 400}},
|
|
]
|
|
workflow = json.dumps({
|
|
"nodes": nodes,
|
|
"links": [
|
|
[1, "1", 0, "2", 0, "MODEL"],
|
|
[2, "2", 0, "3", 0, "IMAGE"],
|
|
[3, "4", 0, "3", 1, "IMAGE"],
|
|
],
|
|
"groups": [{"title": "Subgraph Host"}],
|
|
"extra": {"groupNodes": {"Demo": {"version": 1}}},
|
|
})
|
|
|
|
result, meta = truncate_workflow_smart(workflow, error_node_id="2", max_chars=900)
|
|
|
|
self.assertIn(meta["truncation_method"], {"property_pruning", "node_removal"})
|
|
parsed = json.loads(result)
|
|
self.assertIn("groups", parsed)
|
|
self.assertIn("extra", parsed)
|
|
self.assertEqual(parsed["groups"][0]["title"], "Subgraph Host")
|
|
kept_ids = {str(node["id"]) for node in parsed["nodes"]}
|
|
for link in parsed["links"]:
|
|
self.assertIn(str(link[1]), kept_ids)
|
|
self.assertIn(str(link[3]), kept_ids)
|
|
|
|
if meta["truncation_method"] == "node_removal":
|
|
self.assertNotIn("4", kept_ids)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main(verbosity=2)
|