diff --git a/runpod/endpoint/asyncio/asyncio_runner.py b/runpod/endpoint/asyncio/asyncio_runner.py index cb75144d1..1a960cfd3 100644 --- a/runpod/endpoint/asyncio/asyncio_runner.py +++ b/runpod/endpoint/asyncio/asyncio_runner.py @@ -81,7 +81,7 @@ async def output(self, timeout: int = 0) -> Any: return self.job_output try: - await asyncio.wait_for(self._wait_for_completion(), timeout) + await asyncio.wait_for(self._wait_for_completion(), timeout or None) except asyncio.TimeoutError as exc: raise TimeoutError("Job timed out.") from exc diff --git a/tests/test_endpoint/test_asyncio_runner.py b/tests/test_endpoint/test_asyncio_runner.py index 77237956a..48bb9ae76 100644 --- a/tests/test_endpoint/test_asyncio_runner.py +++ b/tests/test_endpoint/test_asyncio_runner.py @@ -78,6 +78,26 @@ async def json_side_effect(): assert output == "OUTPUT" assert await job.output() == "OUTPUT" + async def test_output_without_timeout_waits_for_completion(self): + """A fresh job can wait for output without an explicit timeout.""" + response = AsyncMock() + response.json.side_effect = [ + {"status": "IN_PROGRESS"}, + {"status": "COMPLETED", "output": "OUTPUT"}, + {"status": "COMPLETED", "output": "OUTPUT"}, + ] + session = AsyncMock() + session.get.return_value = response + job = Job("endpoint_id", "job_id", session, {}) + + with patch( + "runpod.endpoint.asyncio.asyncio_runner.asyncio.sleep", + new_callable=AsyncMock, + ): + output = await job.output() + + self.assertEqual(output, "OUTPUT") + async def test_output_timeout(self): """ Tests Job.output with a timeout