Tool_Agent/tools/local/jimeng_ai/tests/test_jimeng.py

"""Jimeng AI Tool Test Script"""
import asyncio
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from jimeng_client import JimengClient


async def test_text2image():
    """Test text-to-image functionality"""
    print("\n=== Testing Text-to-Image ===")
    
    client = JimengClient(
        api_key=os.getenv("JIMENG_API_KEY"),
        cookie=os.getenv("JIMENG_COOKIE"),
        base_url=os.getenv("JIMENG_BASE_URL", "https://api.jimeng.ai")
    )
    
    try:
        # Submit text-to-image task
        result = await client.text2image(
            prompt="A cute cat playing in the garden",
            negative_prompt="blurry, low quality",
            aspect_ratio="1:1",
            image_count=1,
            cfg_scale=7.0,
            steps=20,
            seed=-1
        )
        
        print(f"Task ID: {result['task_id']}")
        print(f"Status: {result['status']}")
        print(f"Estimated time: {result.get('estimated_time', 'N/A')} seconds")
        
        if result['status'] == 'failed':
            print(f"Error: {result.get('error')}")
            return False
        
        # Simulate query task status
        print("\nQuerying task status...")
        status = await client.query_status(result['task_id'])
        print(f"Current status: {status['status']}")
        print(f"Progress: {status.get('progress', 0)}%")
        
        return True
        
    except Exception as e:
        print(f"Test failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        await client.close()


async def test_image2video():
    """Test image-to-video functionality"""
    print("\n=== Testing Image-to-Video ===")
    
    client = JimengClient(
        api_key=os.getenv("JIMENG_API_KEY"),
        cookie=os.getenv("JIMENG_COOKIE"),
        base_url=os.getenv("JIMENG_BASE_URL", "https://api.jimeng.ai")
    )
    
    try:
        # Use test image URL
        test_image_url = "https://example.com/test.jpg"
        
        result = await client.image2video(
            image_url=test_image_url,
            prompt="Make the image move",
            video_duration=5,
            motion_strength=0.5,
            seed=-1
        )
        
        print(f"Task ID: {result['task_id']}")
        print(f"Status: {result['status']}")
        print(f"Estimated time: {result.get('estimated_time', 'N/A')} seconds")
        
        if result['status'] == 'failed':
            print(f"Error: {result.get('error')}")
            # This is expected since we're using a test URL
            return True
        
        return True
        
    except Exception as e:
        print(f"Test failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        await client.close()


async def test_cache():
    """Test cache functionality"""
    print("\n=== Testing Cache ===")
    
    client = JimengClient()
    
    try:
        # Test cache set and get
        test_data = {
            "task_id": "test_123",
            "status": "completed",
            "result": {"images": ["test.jpg"]}
        }
        
        client.cache.set("test_123", test_data)
        cached = client.cache.get("test_123")
        
        if cached and cached["task_id"] == "test_123":
            print("OK: Cache set and get working")
        else:
            print("FAIL: Cache not working")
            return False
        
        # Test cache cleanup
        client.cleanup_cache()
        print("OK: Cache cleanup working")
        
        return True
        
    except Exception as e:
        print(f"Test failed: {str(e)}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        await client.close()


async def main():
    """Run all tests"""
    print("=" * 50)
    print("Jimeng AI Tool Tests")
    print("=" * 50)
    
    results = []
    
    # Test cache (doesn't need API key)
    results.append(("Cache", await test_cache()))
    
    # Check if API key is configured
    if not os.getenv("JIMENG_API_KEY") and not os.getenv("JIMENG_COOKIE"):
        print("\nWARNING: JIMENG_API_KEY or JIMENG_COOKIE not configured")
        print("Skipping API tests. Configure credentials in .env for full testing")
    else:
        # Test text-to-image
        results.append(("Text2Image", await test_text2image()))
        
        # Test image-to-video
        results.append(("Image2Video", await test_image2video()))
    
    # Output test results
    print("\n" + "=" * 50)
    print("Test Results Summary")
    print("=" * 50)
    
    for name, success in results:
        status = "PASS" if success else "FAIL"
        print(f"{name}: {status}")
    
    all_passed = all(success for _, success in results)
    print("\nOverall:", "ALL TESTS PASSED" if all_passed else "SOME TESTS FAILED")
    
    return all_passed


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)