diff --git a/src/apify_client/_status_message_watcher.py b/src/apify_client/_status_message_watcher.py index 138c2207..d1ccd8d2 100644 --- a/src/apify_client/_status_message_watcher.py +++ b/src/apify_client/_status_message_watcher.py @@ -115,7 +115,8 @@ async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None ) -> None: """Cancel the logging task.""" - await asyncio.sleep(self._final_sleep_time_s) + if exc_type is None: + await asyncio.sleep(self._final_sleep_time_s) await self.stop() async def _log_changed_status_message(self) -> None: @@ -169,7 +170,6 @@ def stop(self) -> None: """Signal the logging thread to stop logging and wait for it to finish.""" if not self._logging_thread: raise RuntimeError('Logging thread is not active') - time.sleep(self._final_sleep_time_s) self._stop_logging = True self._logging_thread.join() self._logging_thread = None @@ -184,6 +184,8 @@ def __exit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None ) -> None: """Stop the logging thread.""" + if exc_type is None: + time.sleep(self._final_sleep_time_s) self.stop() def _log_changed_status_message(self) -> None: diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 492006e0..cf537795 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -44,6 +44,10 @@ ) _EXISTING_LOGS_BEFORE_REDIRECT_ATTACH = 3 +# Large enough that a real sleep is clearly detectable against `_FAST_EXIT_THRESHOLD_S`. +_PATCHED_FINAL_SLEEP_S = 5 +_FAST_EXIT_THRESHOLD_S = 1.0 + _EXPECTED_MESSAGES_AND_LEVELS = ( ('2025-05-13T07:24:12.588Z ACTOR: Pulling Docker image of build.', logging.INFO), ('2025-05-13T07:24:12.686Z ACTOR: Creating Docker container.', logging.INFO), @@ -652,3 +656,64 @@ async def test_status_message_watcher_async_restart_after_normal_completion(http assert task2 is not task # New task created await task2 # Let it complete (will hit terminal status again) assert task2.done() + + +@pytest.mark.usefixtures('mock_api') +def test_sync_watcher_manual_stop_skips_final_sleep( + httpserver: HTTPServer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Manual `stop()` on the sync watcher must not pay the final sleep — only `__exit__` should.""" + monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S) + + api_url = httpserver.url_for('/').removesuffix('/') + run_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID) + watcher = run_client.get_status_message_watcher(check_period=timedelta(seconds=0)) + + watcher.start() + start = time.monotonic() + watcher.stop() + elapsed = time.monotonic() - start + + assert elapsed < _FAST_EXIT_THRESHOLD_S, f'stop() should not sleep, took {elapsed:.2f}s' + + +@pytest.mark.usefixtures('mock_api') +def test_sync_watcher_exit_skips_final_sleep_on_exception( + httpserver: HTTPServer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exceptional `with`-exit must not pay the final sleep so exceptions propagate immediately.""" + monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S) + + api_url = httpserver.url_for('/').removesuffix('/') + run_client = ApifyClient(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID) + watcher = run_client.get_status_message_watcher(check_period=timedelta(seconds=0)) + + start = time.monotonic() + with pytest.raises(RuntimeError, match='boom'), watcher: + raise RuntimeError('boom') + elapsed = time.monotonic() - start + + assert elapsed < _FAST_EXIT_THRESHOLD_S, f'__exit__ should skip final sleep on exception, took {elapsed:.2f}s' + + +@pytest.mark.usefixtures('mock_api') +async def test_async_watcher_aexit_skips_final_sleep_on_exception( + httpserver: HTTPServer, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Exceptional `async with`-exit must not pay the final sleep so exceptions propagate immediately.""" + monkeypatch.setattr(StatusMessageWatcherBase, '_final_sleep_time_s', _PATCHED_FINAL_SLEEP_S) + + api_url = httpserver.url_for('/').removesuffix('/') + run_client = ApifyClientAsync(token='mocked_token', api_url=api_url).run(run_id=_MOCKED_RUN_ID) + watcher = await run_client.get_status_message_watcher(check_period=timedelta(seconds=0)) + + start = time.monotonic() + with pytest.raises(RuntimeError, match='boom'): + async with watcher: + raise RuntimeError('boom') + elapsed = time.monotonic() - start + + assert elapsed < _FAST_EXIT_THRESHOLD_S, f'__aexit__ should skip final sleep on exception, took {elapsed:.2f}s'