test_dag_integration.py

#!/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)