Repository: Significant-Gravitas/Auto-GPT
File: autogpt_platform/backend/backend/util/service_test.py (Python)

Code before:
        self.request_log.append(f"cleanup started at {time.time()}")

        # Call parent cleanup
        super().cleanup()

        self.cleanup_completed = True
        self.request_log.append(f"cleanup completed at {time.time()}")


def run_test_graceful_shutdown_service():
    """Run the test service in a separate process"""

    service = TestGracefulShutdownService()
    service.start()


async def send_slow_request(base_url: str) -> dict:
    """Send a slow request and return the result"""
    async with httpx.AsyncClient(timeout=30) as client:
        response = await client.post(f"{base_url}/slow_endpoint", json={"duration": 5})
        assert response.status_code == 200
        return response.json()


@pytest.mark.asyncio
async def test_graceful_shutdown():
    """Test that AppService handles graceful shutdown correctly"""

    # Start the service in a separate process
    service_process = Process(target=run_test_graceful_shutdown_service)
    service_process.start()

    try:
        # Wait for service to start up
        await asyncio.sleep(3)

        base_url = "http://localhost:18999"

        # Verify service is running
        async with httpx.AsyncClient(timeout=30) as client:
            response = await client.get(f"{base_url}/health_check")
            assert response.status_code == 200
            assert response.json() == "OK"

        # Start a slow request that should complete even after shutdown
        slow_task = asyncio.create_task(send_slow_request(base_url))

        # Give the slow request time to start
        await asyncio.sleep(1)

        # Send SIGTERM to the service process

Review comment (performance):
Result:
```
backend/util/service_test.py::test_graceful_shutdown
2025-10-23 13:22:31,860 INFO  [TestGracefulShutdownService] Starting...
2025-10-23 13:22:31,870 INFO  Metrics endpoint exposed at /metrics for TestGracefulShutdownService
2025-10-23 13:22:31,873 INFO  [PID-13534|THREAD-13545|TestGracefulShutdownService|FastAPI server-702ab79f-9e57-43aa-84d6-58b5c43eb356] Running FastAPI server started...
2025-10-23 13:22:31,873 INFO  [TestGracefulShutdownService] Starting RPC server at http://localhost:18999
2025-10-23 13:22:31,888 INFO  Started server process [13534]
2025-10-23 13:22:31,888 INFO  Waiting for application startup.
2025-10-23 13:22:31,889 INFO  Application startup complete.
2025-10-23 13:22:31,892 INFO  Uvicorn running on http://localhost:18999 (Press CTRL+C to quit)
2025-10-23 13:22:32,765 INFO  127.0.0.1:60804 - "GET /health_check HTTP/1.1" 200

  [▶️ POST /slow_endpoint request is made]
  [🛑 SIGTERM is sent to the service process]

2025-10-23 13:22:33,769 INFO  [TestGracefulShutdownService] 🛑 Received SIGTERM (15) - Entering RPC server graceful shutdown
2025-10-23 13:22:33,801 INFO  Shutting down
2025-10-23 13:22:33,902 INFO  Waiting for connections to close. (CTRL+C to force quit)

  [⌛ 4 seconds pass while POST /slow_endpoint is running]

2025-10-23 13:22:37,804 INFO  127.0.0.1:60816 - "POST /slow_endpoint HTTP/1.1" 200

  [✅ all pending requests have finished -> shutdown can continue]

2025-10-23 13:22:37,819 INFO  Waiting for application shutdown.
2025-10-23 13:22:37,819 INFO  [TestGracefulShutdownService] ✅ FastAPI has finished

  [📨 AppService._lifespan sets BaseAppService._halt_event]

2025-10-23 13:22:37,819 INFO  Application shutdown complete.
2025-10-23 13:22:37,819 INFO  Finished server process [13534]
2025-10-23 13:22:37,819 INFO  [PID-13534|THREAD-13545|TestGracefulShutdownService|FastAPI server-702ab79f-9e57-43aa-84d6-58b5c43eb356] Running FastAPI server completed successfully.

  [🛑 Main thread checks _halt_event and exits its loop]

2025-10-23 13:22:41,874 INFO  [TestGracefulShutdownService] 🧹 Running cleanup
2025-10-23 13:22:41,875 INFO  [TestGracefulShutdownService] ✅ Cleanup done
2025-10-23 13:22:42,573 INFO  [TestGracefulShutdownService] 🛑 Shared event loop stopped
2025-10-23 13:22:42,573 INFO  [TestGracefulShutdownService] 🛑 Terminated
Shutdown took 8.12 seconds
Slow request completed in: 5.00 seconds
PASSED
```

Code after:
        # Call parent cleanup
        super().cleanup()

        self.cleanup_completed = True
        self.request_log.append(f"cleanup completed at {time.time()}")


@pytest.fixture(scope="function")
async def test_service():
    """Run the test service in a separate process"""

    service = TestGracefulShutdownService()
    service.start(background=True)

    base_url = f"http://localhost:{service.get_port()}"

    await wait_until_service_ready(base_url)
    yield service, base_url

    service.stop()


async def wait_until_service_ready(base_url: str, timeout: float = 10):
    start_time = time.time()
    while time.time() - start_time <= timeout:
        async with httpx.AsyncClient(timeout=5) as client:
            with contextlib.suppress(httpx.ConnectError):
                response = await client.get(f"{base_url}/health_check", timeout=5)

                if response.status_code == 200 and response.json() == "OK":
                    return

        await asyncio.sleep(0.5)

    raise RuntimeError(f"Service at {base_url} not available after {timeout} seconds")


async def send_slow_request(base_url: str) -> dict:
    """Send a slow request and return the result"""
    async with httpx.AsyncClient(timeout=30) as client:
        response = await client.post(f"{base_url}/slow_endpoint", json={"duration": 5})
        assert response.status_code == 200
        return response.json()


@pytest.mark.asyncio
async def test_graceful_shutdown(test_service):
    """Test that AppService handles graceful shutdown correctly"""
    service, test_service_url = test_service

    # Start a slow request that should complete even after shutdown