crewai/test_import.py

"""
Quick Test Script
Verifies all modules can be imported and initialized correctly
"""

import sys


def test_imports():
    """Test all module imports"""
    print("=" * 60)
    print("Testing Module Imports...")
    print("=" * 60)
    
    try:
        from agents_config import create_agents, TASK_TEMPLATES, QWEN_MODEL_CONFIG
        print("[OK] agents_config")
    except Exception as e:
        print(f"[FAIL] agents_config - {e}")
        return False
    
    try:
        from stream_manager import stream_manager, StreamEvent, TaskStreamQueue
        print("[OK] stream_manager")
    except Exception as e:
        print(f"[FAIL] stream_manager - {e}")
        return False
    
    try:
        from crew_factory import CrewFactory, SSECrewExecutor, CrewExecutionLogger
        print("[OK] crew_factory")
    except Exception as e:
        print(f"[FAIL] crew_factory - {e}")
        return False
    
    try:
        from main import app
        print("[OK] main")
    except Exception as e:
        print(f"[FAIL] main - {e}")
        return False
    
    return True


def test_agents():
    """Test Agent creation"""
    print("\n" + "=" * 60)
    print("Testing Agent Creation...")
    print("=" * 60)
    
    try:
        from agents_config import create_agents, get_product_manager_agent
        # Test creating a single agent without LLM configuration
        agent = get_product_manager_agent()
        
        print(f"[OK] Agent structure created:")
        print(f"  - Role: {agent.role}")
        print(f"  - Goal: {agent.goal[:50]}...")
        print(f"\n[NOTE] Full agent initialization requires DASHSCOPE_API_KEY")
        print(f"       Set environment variable before running the server.")
        
        return True
    except Exception as e:
        # This is expected if API key is not configured
        error_msg = str(e)
        if "API_KEY" in error_msg or "provider" in error_msg.lower():
            print(f"[SKIP] Agent creation skipped (API key not configured)")
            print(f"       Set DASHSCOPE_API_KEY environment variable to enable.")
            return True  # Not a failure, just needs configuration
        else:
            print(f"[FAIL] Agent creation failed: {e}")
            return False


def test_stream_manager():
    """Test stream manager"""
    print("\n" + "=" * 60)
    print("Testing Stream Manager...")
    print("=" * 60)
    
    try:
        import asyncio
        from stream_manager import stream_manager, StreamEvent
        
        async def test():
            # Create test stream
            task_id = "test-123"
            await stream_manager.create_stream(task_id)
            
            # Publish test event
            await stream_manager.publish_event(
                task_id=task_id,
                event_type="test",
                agent="TestAgent",
                content="This is a test event"
            )
            
            # Get stream
            stream = await stream_manager.get_stream(task_id)
            if stream:
                event = await stream.get()
                if event:
                    print(f"[OK] Event format: {event.to_dict()}")
                    await stream_manager.close_stream(task_id)
                    return True
            
            return False
        
        result = asyncio.get_event_loop().run_until_complete(test())
        if result:
            print("[OK] Stream manager test passed")
            return True
        else:
            print("[FAIL] Stream manager test failed")
            return False
            
    except Exception as e:
        print(f"[FAIL] Stream manager test failed: {e}")
        return False


def test_api_endpoints():
    """Test API endpoint registration"""
    print("\n" + "=" * 60)
    print("Testing API Endpoints...")
    print("=" * 60)
    
    try:
        from main import app
        
        routes = [route.path for route in app.routes]
        
        required_endpoints = [
            "/api/run_task",
            "/api/stream/{task_id}",
            "/api/task/{task_id}/status",
            "/api/streams",
            "/health",
            "/test-ui"
        ]
        
        print(f"[OK] Registered endpoints ({len(routes)} total):")
        for endpoint in required_endpoints:
            found = False
            for r in routes:
                if endpoint.split('{')[0].rstrip('/') in r:
                    print(f"  [OK] {endpoint}")
                    found = True
                    break
            if not found:
                print(f"  [?] {endpoint} (may use different format)")
        
        return True
    except Exception as e:
        print(f"[FAIL] API endpoint test failed: {e}")
        return False


def main():
    """Run all tests"""
    print("\n" + "=" * 60)
    print("Multi-Agent System Module Test")
    print("=" * 60 + "\n")
    
    results = []
    
    # Test imports
    results.append(("Module Imports", test_imports()))
    
    # Test Agent creation
    results.append(("Agent Creation", test_agents()))
    
    # Test stream manager
    results.append(("Stream Manager", test_stream_manager()))
    
    # Test API endpoints
    results.append(("API Endpoints", test_api_endpoints()))
    
    # Summary
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    
    for name, passed in results:
        status = "[OK]" if passed else "[FAIL]"
        print(f"{status} - {name}")
    
    all_passed = all(result[1] for result in results)
    
    print("\n" + "=" * 60)
    if all_passed:
        print("[SUCCESS] All tests passed! System is ready.")
        print("\nTo start the server:")
        print("  python main.py")
        print("\nTest UI:")
        print("  http://localhost:8000/test-ui")
        print("\nAPI Documentation:")
        print("  http://localhost:8000/docs")
    else:
        print("[FAILURE] Some tests failed. Please check error messages.")
        sys.exit(1)
    print("=" * 60)


if __name__ == "__main__":
    main()