Review comments

Author: Pijukatel

src/apify_client/clients/resource_clients/log.py

@@ -12,7 +12,6 @@
 from typing import TYPE_CHECKING, Any, cast
 
 from apify_shared.utils import ignore_docs
-from typing_extensions import Self
 
 from apify_client._errors import ApifyApiError
 from apify_client._utils import catch_not_found_or_throw
@@ -354,13 +353,16 @@ def start(self) -> Task:
         self._streaming_task = asyncio.create_task(self._stream_log())
         return self._streaming_task
 
-    def stop(self) -> None:
+    async def stop(self) -> None:
         """Stop the streaming task."""
         if not self._streaming_task:
             raise RuntimeError('Streaming task is not active')
 
         self._streaming_task.cancel()
-        self._streaming_task = None
+        try:
+            await self._streaming_task
+        except asyncio.CancelledError:
+            self._streaming_task = None
 
     async def __aenter__(self) -> Self:
         """Start the streaming task within the context. Exiting the context will cancel the streaming task."""
@@ -371,7 +373,7 @@ async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """Cancel the streaming task."""
-        self.stop()
+        await self.stop()
 
     async def _stream_log(self) -> None:
         async with self._log_client.stream(raw=True) as log_stream:
@@ -456,13 +458,16 @@ def start(self) -> Task:
         self._logging_task = asyncio.create_task(self._log_changed_status_message())
         return self._logging_task
 
-    def stop(self) -> None:
+    async def stop(self) -> None:
         """Stop the logging task."""
         if not self._logging_task:
             raise RuntimeError('Logging task is not active')
 
         self._logging_task.cancel()
-        self._logging_task = None
+        try:
+            await self._logging_task
+        except asyncio.CancelledError:
+            self._logging_task = None
 
     async def __aenter__(self) -> Self:
         """Start the logging task within the context. Exiting the context will cancel the logging task."""
@@ -474,7 +479,7 @@ async def __aexit__(
     ) -> None:
         """Cancel the logging task."""
         await asyncio.sleep(self._final_sleep_time_s)
-        self.stop()
+        await self.stop()
 
     async def _log_changed_status_message(self) -> None:
         while True:

src/apify_client/clients/resource_clients/run.py

@@ -262,7 +262,7 @@ def log(self) -> LogClient:
     def get_streamed_log(self, to_logger: logging.Logger | None = None, *, from_start: bool = True) -> StreamedLogSync:
         """Get `StreamedLog` instance that can be used to redirect logs.
 
-         `StreamedLog` can be directly called or used as a context manager.
+         `StreamedLog` can be explicitly started and stopped or used as a context manager.
 
         Args:
             to_logger: `Logger` used for logging the redirected messages. If not provided, a new logger is created
@@ -327,7 +327,7 @@ def get_status_message_watcher(
     ) -> StatusMessageWatcherSync:
         """Get `StatusMessageWatcher` instance that can be used to redirect status and status messages to logs.
 
-        `StatusMessageWatcher` can be directly called or used as a context manager.
+        `StatusMessageWatcher` can be explicitly started and stopped or used as a context manager.
 
         Args:
             to_logger: `Logger` used for logging the status and status messages. If not provided, a new logger is
@@ -586,7 +586,7 @@ async def get_streamed_log(
     ) -> StreamedLogAsync:
         """Get `StreamedLog` instance that can be used to redirect logs.
 
-         `StreamedLog` can be directly called or used as a context manager.
+         `StreamedLog` can be explicitly started and stopped or used as a context manager.
 
         Args:
             to_logger: `Logger` used for logging the redirected messages. If not provided, a new logger is created
@@ -652,7 +652,7 @@ async def get_status_message_watcher(
     ) -> StatusMessageWatcherAsync:
         """Get `StatusMessageWatcher` instance that can be used to redirect status and status messages to logs.
 
-        `StatusMessageWatcher` can be directly called or used as a context manager.
+        `StatusMessageWatcher` can be explicitly started and stopped or used as a context manager.
 
         Args:
             to_logger: `Logger` used for logging the status and status messages. If not provided, a new logger is

tests/unit/test_logging.py

@@ -2,7 +2,7 @@
 import json
 import logging
 import time
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Generator, Iterator
 from datetime import datetime, timedelta
 from unittest.mock import patch
 
@@ -65,7 +65,7 @@
 
 @pytest.fixture
 def mock_api() -> None:
-    def create_status_responses_generator() -> Iterator[httpx.Response]:
+    def get_responses() -> Generator[httpx.Response, None, None]:
         """Simulate actor run that changes status 3 times."""
         for _ in range(5):
             yield httpx.Response(
@@ -113,11 +113,11 @@ def create_status_responses_generator() -> Iterator[httpx.Response]:
                 status_code=200,
             )
 
-    response_generator = create_status_responses_generator()
+    responses = get_responses()
 
     def actor_runs_side_effect(_: httpx.Request) -> httpx.Response:
         time.sleep(0.1)
-        return next(response_generator)
+        return next(responses)
 
     respx.get(url=f'{_MOCKED_API_URL}/v2/actor-runs/{_MOCKED_RUN_ID}').mock(side_effect=actor_runs_side_effect)