#!/usr/bin/env python3 """ Test script for DAG integration in DART-LLM-Multi-Model """ from dag_visualizer import DAGVisualizer from json_processor import JsonProcessor import json def test_dag_integration(): """Test the DAG integration with sample task data""" print("Testing DAG integration...") # Sample response with task data (similar to what the model might generate) sample_response = """ Based on your command, here are the robot tasks: { "tasks": [ { "task": "move_excavator_to_soil_area", "instruction_function": { "name": "move_to_position", "robot_ids": ["robot_excavator_01"], "dependencies": [], "object_keywords": ["soil_area_1"] } }, { "task": "excavate_soil", "instruction_function": { "name": "excavate_material", "robot_ids": ["robot_excavator_01"], "dependencies": ["move_excavator_to_soil_area"], "object_keywords": ["soil"] } }, { "task": "move_dump_truck", "instruction_function": { "name": "move_to_position", "robot_ids": ["robot_dump_truck_01"], "dependencies": [], "object_keywords": ["soil_area_1"] } }, { "task": "load_soil_to_truck", "instruction_function": { "name": "transfer_material", "robot_ids": ["robot_excavator_01", "robot_dump_truck_01"], "dependencies": ["excavate_soil", "move_dump_truck"], "object_keywords": ["soil"] } }, { "task": "transport_to_dump_site", "instruction_function": { "name": "move_to_position", "robot_ids": ["robot_dump_truck_01"], "dependencies": ["load_soil_to_truck"], "object_keywords": ["dump_site"] } } ] } """ # Test JSON processing processor = JsonProcessor() print("1. Testing JSON processing...") processed_json = processor.process_response(sample_response) if processed_json: print("✓ JSON processing successful") print(f" Found {len(processed_json['tasks'])} tasks") else: print("✗ JSON processing failed") return False # Test DAG visualization print("2. Testing DAG visualization...") visualizer = DAGVisualizer() try: dag_image_path = visualizer.create_dag_visualization( processed_json, title="Test Robot Task Dependency Graph" ) if dag_image_path: print(f"✓ DAG visualization created: {dag_image_path}") return True else: print("✗ DAG visualization failed") return False except Exception as e: print(f"✗ DAG visualization error: {e}") return False def test_simplified_dag(): """Test simplified DAG visualization""" print("\n3. Testing simplified DAG...") simple_task_data = { "tasks": [ { "task": "move_robot", "instruction_function": { "name": "move_to_position", "robot_ids": ["robot_01"], "dependencies": [], "object_keywords": ["target_area"] } }, { "task": "perform_operation", "instruction_function": { "name": "excavate", "robot_ids": ["robot_01"], "dependencies": ["move_robot"], "object_keywords": ["soil"] } } ] } visualizer = DAGVisualizer() try: dag_image_path = visualizer.create_simplified_dag_visualization( simple_task_data, title="Simplified Test DAG" ) if dag_image_path: print(f"✓ Simplified DAG visualization created: {dag_image_path}") return True else: print("✗ Simplified DAG visualization failed") return False except Exception as e: print(f"✗ Simplified DAG visualization error: {e}") return False def main(): """Run all tests""" print("=" * 60) print("DART-LLM-Multi-Model DAG Integration Test") print("=" * 60) success_count = 0 total_tests = 2 if test_dag_integration(): success_count += 1 if test_simplified_dag(): success_count += 1 print("\n" + "=" * 60) print(f"Test Results: {success_count}/{total_tests} passed") if success_count == total_tests: print("🎉 All DAG integration tests passed!") return True else: print("❌ Some tests failed!") return False if __name__ == "__main__": success = main() exit(0 if success else 1)