openenv/atari_env

1#!/usr/bin/env python32"""3End-to-end test for LocalDockerProvider.4 5This script tests the complete flow:61. Start a container using LocalDockerProvider72. Wait for it to be ready83. Make HTTP requests to test the environment94. Clean up the container10"""11 12import sys13from pathlib import Path14 15# Add src to path16sys.path.insert(0, str(Path(__file__).parent.parent.parent))17 18import requests19from openenv.core.containers.runtime import LocalDockerProvider20 21 22# TODO: Remove this test or make it a functional test sicne this will be tested in e2e test for echo env23def test_local_docker_provider():24    """Test LocalDockerProvider end-to-end."""25    print("=" * 60)26    print("LocalDockerProvider End-to-End Test")27    print("=" * 60)28    print()29 30    provider = None31 32    try:33        # Step 1: Create provider34        print("Step 1: Creating LocalDockerProvider...")35        provider = LocalDockerProvider()36        print("✓ Provider created\n")37 38        # Step 2: Start container39        print("Step 2: Starting echo-env container...")40        base_url = provider.start_container("echo-env:latest")41        print(f"✓ Container started at: {base_url}")42        if provider._container_id:43            print(f"  Container ID: {provider._container_id[:12]}...")44        if provider._container_name:45            print(f"  Container name: {provider._container_name}\n")46 47        # Step 3: Wait for ready48        print("Step 3: Waiting for container to be ready...")49        provider.wait_for_ready(base_url, timeout_s=30.0)50        print("✓ Container is ready!\n")51 52        # Step 4: Test health endpoint53        print("Step 4: Testing /health endpoint...")54        response = requests.get(f"{base_url}/health")55        print(f"  Status: {response.status_code}")56        print(f"  Response: {response.json()}")57        assert response.status_code == 20058        assert response.json()["status"] == "healthy"59        print("✓ Health check passed\n")60 61        # Step 5: Test reset endpoint62        print("Step 5: Testing /reset endpoint...")63        response = requests.post(64            f"{base_url}/reset",65            json={},66            headers={"Content-Type": "application/json"},67        )68        print(f"  Status: {response.status_code}")69        data = response.json()70        print(f"  Message: {data['observation']['echoed_message']}")71        print(f"  Reward: {data['reward']}")72        print(f"  Done: {data['done']}")73        assert response.status_code == 20074        assert data["observation"]["echoed_message"] == "Echo environment ready!"75        print("✓ Reset test passed\n")76 77        # Step 6: Test step endpoint78        print("Step 6: Testing /step endpoint...")79        response = requests.post(80            f"{base_url}/step",81            json={"action": {"message": "Hello from LocalDockerProvider!"}},82            headers={"Content-Type": "application/json"},83        )84        print(f"  Status: {response.status_code}")85        data = response.json()86        print(f"  Echoed: {data['observation']['echoed_message']}")87        print(f"  Length: {data['observation']['message_length']}")88        print(f"  Reward: {data['reward']}")89        assert response.status_code == 20090        assert (91            data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"92        )93        assert data["observation"]["message_length"] == 3194        print("✓ Step test passed\n")95 96        # Step 7: Test state endpoint97        print("Step 7: Testing /state endpoint...")98        response = requests.get(f"{base_url}/state")99        print(f"  Status: {response.status_code}")100        data = response.json()101        print(f"  Episode ID: {data['episode_id']}")102        print(f"  Step count: {data['step_count']}")103        assert response.status_code == 200104        assert data["step_count"] == 1  # One step from above105        print("✓ State test passed\n")106 107        # Step 8: Multiple steps108        print("Step 8: Testing multiple steps...")109        for i in range(3):110            response = requests.post(111                f"{base_url}/step",112                json={"action": {"message": f"Message {i + 1}"}},113                headers={"Content-Type": "application/json"},114            )115            assert response.status_code == 200116            print(f"  Step {i + 1}: ✓")117 118        # Check state updated119        response = requests.get(f"{base_url}/state")120        data = response.json()121        assert data["step_count"] == 4  # 1 + 3 more steps122        print(f"  Final step count: {data['step_count']}")123        print("✓ Multiple steps test passed\n")124 125        print("=" * 60)126        print("✓ All tests passed!")127        print("=" * 60)128        print()129 130        return True131 132    except Exception as e:133        print(f"\n❌ Test failed: {e}")134        import traceback135 136        traceback.print_exc()137        return False138 139    finally:140        # Step 9: Cleanup141        if provider is not None:142            print("\nStep 9: Cleaning up container...")143            try:144                provider.stop_container()145                print("✓ Container stopped and removed\n")146            except Exception as e:147                print(f"⚠️  Cleanup warning: {e}\n")148 149 150def test_provider_with_custom_port():151    """Test provider with custom port."""152    print("=" * 60)153    print("LocalDockerProvider with Custom Port Test")154    print("=" * 60)155    print()156 157    provider = None158 159    try:160        provider = LocalDockerProvider()161 162        print("Starting container on custom port 8123...")163        base_url = provider.start_container("echo-env:latest", port=8123)164        print(f"✓ Started at: {base_url}")165        assert ":8123" in base_url166 167        print("Waiting for ready...")168        provider.wait_for_ready(base_url)169        print("✓ Ready!")170 171        print("Testing health...")172        response = requests.get(f"{base_url}/health")173        assert response.status_code == 200174        print("✓ Health check passed")175 176        print("\n✓ Custom port test passed!\n")177        return True178 179    except Exception as e:180        print(f"\n❌ Test failed: {e}")181        return False182 183    finally:184        if provider is not None:185            provider.stop_container()186            print("✓ Cleaned up\n")187 188 189def test_provider_with_env_vars():190    """Test provider with environment variables."""191    print("=" * 60)192    print("LocalDockerProvider with Environment Variables Test")193    print("=" * 60)194    print()195 196    provider = None197 198    try:199        provider = LocalDockerProvider()200 201        print("Starting container with environment variables...")202        base_url = provider.start_container(203            "echo-env:latest", env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}204        )205        print(f"✓ Started at: {base_url}")206 207        print("Waiting for ready...")208        provider.wait_for_ready(base_url)209        print("✓ Ready!")210 211        print("Testing health...")212        response = requests.get(f"{base_url}/health")213        assert response.status_code == 200214        print("✓ Health check passed")215 216        print("\n✓ Environment variables test passed!\n")217        return True218 219    except Exception as e:220        print(f"\n❌ Test failed: {e}")221        return False222 223    finally:224        if provider is not None:225            provider.stop_container()226            print("✓ Cleaned up\n")227 228 229if __name__ == "__main__":230    print()231    print("🐳 LocalDockerProvider Test Suite")232    print()233 234    results = []235 236    # Run basic test237    results.append(("Basic End-to-End", test_local_docker_provider()))238 239    # Run custom port test240    results.append(("Custom Port", test_provider_with_custom_port()))241 242    # Run environment variables test243    results.append(("Environment Variables", test_provider_with_env_vars()))244 245    # Summary246    print("=" * 60)247    print("Test Summary")248    print("=" * 60)249    for name, passed in results:250        status = "✓ PASSED" if passed else "✗ FAILED"251        print(f"{name:25} {status}")252    print("=" * 60)253 254    all_passed = all(result for _, result in results)255    if all_passed:256        print("\n🎉 All tests passed!")257        exit(0)258    else:259        print("\n❌ Some tests failed")260        exit(1)261