from mjr_am_backend.features.workflows import classify_workflow, parse_workflow, workflow_node_text def test_parse_workflow_includes_nested_subgraph_nodes(): subgraph = { "nodes": [ {"id": 10, "type": "LoadImage", "widgets_values": ["input.png"]}, {"id": 11, "type": "KSampler"}, ], "links": [[1, 10, 0, 11, 0, "IMAGE"]], } workflow = { "nodes": [ {"id": 1, "type": "CheckpointLoaderSimple", "widgets_values": ["model.safetensors"]}, {"id": 2, "type": "GroupNode", "subgraph": subgraph}, ], "links": [[1, 1, 0, 2, 0, "MODEL"]], } stats = parse_workflow(workflow) assert stats.node_count == 2 assert stats.link_count == 1 assert stats.subgraph_count == 1 assert stats.subgraph_node_count == 2 assert any(node.get("id") == 10 for node in stats.nodes) assert any(node.get("id") == 11 for node in stats.nodes) assert "root::1" in (stats.qualified_node_ids or []) assert any(item.endswith("::10") for item in (stats.qualified_node_ids or [])) def test_workflow_node_text_collects_widget_and_input_text(): text = workflow_node_text( [ { "type": "CLIPTextEncode", "widgets_values": ["cinematic lighting", 123], "inputs": {"text": "prompt text", "strength": 0.8}, } ] ) assert "cliptextencode" in text assert "cinematic lighting" in text assert "prompt text" in text def test_classify_workflow_detects_common_workflow_buckets(): video_edit = classify_workflow("LoadVideo Wan cloud provider") i2v = classify_workflow("LoadImage Seedance workflow") upscale = classify_workflow("Ultimate upscale model") image_edit = classify_workflow("inpaint image edit") assert video_edit.task == "Video Edit" assert video_edit.runs_on == "api" assert i2v.task == "I2V" assert i2v.model_family == "Seedance" assert upscale.task == "Upscale" assert image_edit.task == "Image Edit" def test_classify_workflow_uses_graph_structure_for_i2v(): result = classify_workflow( "", [ {"id": 1, "type": "LoadImage", "widgets_values": ["start.png"]}, {"id": 2, "type": "WanVideoSampler"}, {"id": 3, "type": "VHS_VideoCombine"}, ], ) assert result.task == "I2V" assert result.model_family == "Wan" assert result.source == "graph" assert result.confidence >= 0.9 assert result.signals["has_image_input"] is True assert result.signals["has_video_sink"] is True def test_classify_workflow_detects_video_edit_and_api_provider(): result = classify_workflow( "", [ {"id": 1, "type": "LoadVideo", "widgets_values": ["clip.mp4"]}, {"id": 2, "type": "GoogleVeoVideoGeneration"}, ], ) assert result.task == "Video Edit" assert result.model_family == "Veo" assert result.provider == "Google" assert result.runs_on == "api" def test_classify_workflow_detects_model_from_checkpoint_widget(): result = classify_workflow( "", [ { "id": 1, "type": "CheckpointLoaderSimple", "widgets_values": ["qwen_image_edit_fp8.safetensors"], }, {"id": 2, "type": "KSampler"}, {"id": 3, "type": "SaveImage"}, ], ) assert result.task == "T2I" assert result.model_family == "Qwen" assert result.runs_on == "local"