Add support for terminal status message, allow setting status message directly in exit (#88)

Depends on apify/apify-client-python#122

---------

Co-authored-by: František Nesveda <[email protected]>

Author: mvolfik

CHANGELOG.md

@@ -7,6 +7,8 @@ Changelog
 ### Added
 
 - option to add event handlers which accept no arguments
+- added support for `is_terminal` flag in status message update
+- option to set status message along with `Actor.exit()`
 
 ### Fixed
 

docs/docs.md

@@ -338,7 +338,7 @@ and it should be called only once.
 
 ***
 
-#### [](#actor-exit) `async Actor.exit(*, exit_code=0, event_listeners_timeout_secs=5)`
+#### [](#actor-exit) `async Actor.exit(*, exit_code=0, event_listeners_timeout_secs=5, status_message=None)`
 
 Exit the actor instance.
 
@@ -352,15 +352,17 @@ and stops the event manager.
 
   * **exit_code** (`int`, *optional*) – The exit code with which the actor should fail (defaults to 0).
 
-  * **event_listeners_timeout_secs** (`float`, *optional*) – How long should the actor wait for actor event listeners to finish before exiting
+  * **event_listeners_timeout_secs** (`float`, *optional*) – How long should the actor wait for actor event listeners to finish before exiting.
+
+  * **status_message** (`str`, *optional*) – The final status message that the actor should display.
 
 * **Return type**
 
   `None`
 
 ***
 
-#### [](#actor-fail) `async Actor.fail(*, exit_code=1, exception=None)`
+#### [](#actor-fail) `async Actor.fail(*, exit_code=1, exception=None, status_message=None)`
 
 Fail the actor instance.
 
@@ -373,6 +375,8 @@ but it additionally sets the exit code to 1 (by default).
 
   * **exception** (`BaseException`, *optional*) – The exception with which the actor failed.
 
+  * **status_message** (`str`, *optional*) – The final status message that the actor should display.
+
 * **Return type**
 
   `None`
@@ -908,14 +912,16 @@ For more information about Apify actor webhooks, please see the [documentation](
 
 ***
 
-#### [](#actor-set_status_message) `async Actor.set_status_message(status_message)`
+#### [](#actor-set_status_message) `async Actor.set_status_message(status_message, *, is_terminal=None)`
 
 Set the status message for the current actor run.
 
 * **Parameters**
 
   * **status_message** (`str`) – The status message to set to the run.
 
+  * **is_terminal** (`bool`, *optional*) – Set this flag to True if this is the final status message of the Actor run.
+
 * **Returns**
 
   The updated actor run object

setup.py

@@ -56,7 +56,7 @@
     install_requires=[
         'aiofiles ~= 22.1.0',
         'aioshutil ~= 1.2',
-        'apify-client ~= 1.0.0',
+        'apify-client ~= 1.1',
         'colorama ~= 0.4.6',
         'cryptography ~= 39.0.1',
         'httpx ~= 0.23.0',

src/apify/actor.py

@@ -301,6 +301,7 @@ async def exit(
         *,
         exit_code: int = 0,
         event_listeners_timeout_secs: Optional[float] = EVENT_LISTENERS_TIMEOUT_SECS,
+        status_message: Optional[str] = None,
     ) -> None:
         """Exit the actor instance.
 
@@ -312,18 +313,21 @@ async def exit(
 
         Args:
             exit_code (int, optional): The exit code with which the actor should fail (defaults to `0`).
-            event_listeners_timeout_secs (float, optional): How long should the actor wait for actor event listeners to finish before exiting
+            event_listeners_timeout_secs (float, optional): How long should the actor wait for actor event listeners to finish before exiting.
+            status_message (str, optional): The final status message that the actor should display.
         """
         return await cls._get_default_instance().exit(
             exit_code=exit_code,
             event_listeners_timeout_secs=event_listeners_timeout_secs,
+            status_message=status_message,
         )
 
     async def _exit_internal(
         self,
         *,
         exit_code: int = 0,
         event_listeners_timeout_secs: Optional[float] = EVENT_LISTENERS_TIMEOUT_SECS,
+        status_message: Optional[str] = None,
     ) -> None:
         self._raise_if_not_initialized()
 
@@ -340,6 +344,9 @@ async def _exit_internal(
             self._event_manager.emit(ActorEventTypes.PERSIST_STATE, {'isMigrating': False})
             self._was_final_persist_state_emitted = True
 
+        if status_message is not None:
+            await self.set_status_message(status_message, is_terminal=True)
+
         # Sleep for a bit so that the listeners have a chance to trigger
         await asyncio.sleep(0.1)
 
@@ -362,6 +369,7 @@ async def fail(
         *,
         exit_code: int = 1,
         exception: Optional[BaseException] = None,
+        status_message: Optional[str] = None,
     ) -> None:
         """Fail the actor instance.
 
@@ -371,17 +379,20 @@ async def fail(
         Args:
             exit_code (int, optional): The exit code with which the actor should fail (defaults to `1`).
             exception (BaseException, optional): The exception with which the actor failed.
+            status_message (str, optional): The final status message that the actor should display.
         """
         return await cls._get_default_instance().fail(
             exit_code=exit_code,
             exception=exception,
+            status_message=status_message,
         )
 
     async def _fail_internal(
         self,
         *,
         exit_code: int = 1,
         exception: Optional[BaseException] = None,
+        status_message: Optional[str] = None,
     ) -> None:
         self._raise_if_not_initialized()
 
@@ -390,7 +401,7 @@ async def _fail_internal(
         if exception and not _is_running_in_ipython():
             self.log.exception('Actor failed with an exception', exc_info=exception)
 
-        await self.exit(exit_code=exit_code)
+        await self.exit(exit_code=exit_code, status_message=status_message)
 
     @classmethod
     async def main(cls, main_actor_function: Callable[[], MainReturnType]) -> Optional[MainReturnType]:
@@ -1210,28 +1221,30 @@ async def _add_webhook_internal(
         )
 
     @classmethod
-    async def set_status_message(cls, status_message: str) -> Optional[Dict]:
+    async def set_status_message(cls, status_message: str, *, is_terminal: Optional[bool] = None) -> Optional[Dict]:
         """Set the status message for the current actor run.
 
         Args:
             status_message (str): The status message to set to the run.
+            is_terminal (bool, optional): Set this flag to True if this is the final status message of the Actor run.
 
         Returns:
             dict: The updated actor run object
         """
-        return await cls._get_default_instance().set_status_message(status_message=status_message)
+        return await cls._get_default_instance().set_status_message(status_message=status_message, is_terminal=is_terminal)
 
-    async def _set_status_message_internal(self, status_message: str) -> Optional[Dict]:
+    async def _set_status_message_internal(self, status_message: str, *, is_terminal: Optional[bool] = None) -> Optional[Dict]:
         self._raise_if_not_initialized()
 
         if not self.is_at_home():
-            self.log.error('Actor.set_status_message() is only supported when running on the Apify platform.')
+            title = 'Terminal status message' if is_terminal else 'Status message'
+            self.log.info(f'[{title}]: {status_message}')
             return None
 
         # If is_at_home() is True, config.actor_run_id is always set
         assert self._config.actor_run_id is not None
 
-        return await self._apify_client.run(self._config.actor_run_id).update(status_message=status_message)
+        return await self._apify_client.run(self._config.actor_run_id).update(status_message=status_message, is_status_message_terminal=is_terminal)
 
     @classmethod
     async def create_proxy_configuration(

tests/integration/test_actor_api_helpers.py

@@ -81,7 +81,8 @@ class TestActorSetStatusMessage:
     async def test_actor_set_status_message(self, make_actor: ActorFactory) -> None:
         async def main() -> None:
             async with Actor:
-                await Actor.set_status_message('testing-status-message')
+                input = await Actor.get_input() or {}
+                await Actor.set_status_message('testing-status-message', **input)
 
         actor = await make_actor('set-status-message', main_func=main)
 
@@ -90,6 +91,14 @@ async def main() -> None:
         assert run_result is not None
         assert run_result['status'] == 'SUCCEEDED'
         assert run_result['statusMessage'] == 'testing-status-message'
+        assert run_result['isStatusMessageTerminal'] is None
+
+        run_result = await actor.call(run_input={'is_terminal': True})
+
+        assert run_result is not None
+        assert run_result['status'] == 'SUCCEEDED'
+        assert run_result['statusMessage'] == 'testing-status-message'
+        assert run_result['isStatusMessageTerminal'] is True
 
 
 class TestActorStart:

tests/integration/test_actor_lifecycle.py

@@ -85,6 +85,12 @@ async def main() -> None:
             assert run_result['exitCode'] == exit_code
             assert run_result['status'] == 'FAILED'
 
+        # fail with status message
+        run_result = await actor.call(run_input={'status_message': 'This is a test message'})
+        assert run_result is not None
+        assert run_result['status'] == 'FAILED'
+        assert run_result.get('statusMessage') == 'This is a test message'
+
     async def test_with_actor_fail_correctly(self, make_actor: ActorFactory) -> None:
         async def main() -> None:
             async with Actor:

tests/unit/actor/test_actor_helpers.py

@@ -102,9 +102,20 @@ async def test_actor_add_webhook_not_work_locally(self, caplog: pytest.LogCaptur
         assert caplog.records[0].levelname == 'ERROR'
         assert 'Actor.add_webhook() is only supported when running on the Apify platform.' in caplog.records[0].message
 
-    async def test_actor_set_status_message_not_work_locally(self, caplog: pytest.LogCaptureFixture) -> None:
+    async def test_actor_set_status_message_mock_locally(self, caplog: pytest.LogCaptureFixture) -> None:
+        caplog.set_level('INFO')
         async with Actor() as my_actor:
-            await my_actor.set_status_message('test')
-        assert len(caplog.records) == 1
-        assert caplog.records[0].levelname == 'ERROR'
-        assert 'Actor.set_status_message() is only supported when running on the Apify platform.' in caplog.records[0].message
+            await my_actor.set_status_message('test-status-message')
+        matching_records = [record for record in caplog.records if 'test-status-message' in record.message]
+        assert len(matching_records) == 1
+        assert matching_records[0].levelname == 'INFO'
+        assert '[Status message]: test-status-message' in matching_records[0].message
+
+    async def test_actor_set_status_message_terminal_mock_locally(self, caplog: pytest.LogCaptureFixture) -> None:
+        caplog.set_level('INFO')
+        async with Actor() as my_actor:
+            await my_actor.fail(status_message='test-terminal-message')
+        matching_records = [record for record in caplog.records if 'test-terminal-message' in record.message]
+        assert len(matching_records) == 1
+        assert matching_records[0].levelname == 'INFO'
+        assert '[Terminal status message]: test-terminal-message' in matching_records[0].message

tests/unit/actor/test_actor_log.py

@@ -5,12 +5,12 @@
 
 from apify import Actor, __version__
 from apify.log import logger
-from apify_client._version import __version__ as apify_client_version
+from apify_client import __version__ as apify_client_version
 
 
 class TestActorLog:
     async def test_actor_log(self, caplog: pytest.LogCaptureFixture) -> None:
-        caplog.set_level(logging.DEBUG)
+        caplog.set_level(logging.DEBUG, logger='apify')
         try:
             async with Actor:
                 # Test Actor.log

tests/unit/test_event_manager.py

@@ -17,7 +17,7 @@
 
 class TestEventManagerLocal:
     async def test_lifecycle_local(self, caplog: pytest.LogCaptureFixture) -> None:
-        caplog.set_level(logging.DEBUG)
+        caplog.set_level(logging.DEBUG, logger='apify')
 
         config = Configuration()
         event_manager = EventManager(config)