Repository: hikari-py/hikari
Pull request #999 (2022-02-03 17:39:24)

Add startup and shutdown callbacks to the RESTBot interface
### Summary
Add startup and shutdown callbacks to the RESTBot interface

### Checklist
<!-- Make sure to tick all the following boxes by putting an `x` in between (like this `[x]`) -->
- [x] I have run `nox` and all the pipelines have passed.
- [x] I have made unittests according to the code I have added/modified/deleted.

### Related issues
<!--
To mention an issue use `#issue-id` and to mention a merge request use `!merge-request-id`
To close/fix an issue use `Close #issue-id` or `Fix #issue-id` (depending on the merge request)
-->



Patch:
diff --git a/MANIFEST.in b/MANIFEST.in
index 10128ec7..4d7cba72 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -2,5 +2,8 @@ graft hikari
 include LICENSE
 include setup.py
 include README.md
+include requirements.txt
+include server-requirements.txt
+inlcude speedup-requirements.txt
 include pyproject.toml
 global-exclude *.py[cod]
diff --git a/changes/999.feature.md b/changes/999.feature.md
new file mode 100644
index 00000000..88e1013e
--- /dev/null
+++ b/changes/999.feature.md
@@ -0,0 +1,1 @@
+Startup and shutdown callbacks for the RESTBot interface/impl.
diff --git a/hikari/impl/rest_bot.py b/hikari/impl/rest_bot.py
index 776fb6ca..abf63330 100644
--- a/hikari/impl/rest_bot.py
+++ b/hikari/impl/rest_bot.py
@@ -182,11 +182,13 @@ class RESTBot(traits.RESTBotAware, interaction_server_.InteractionServer):
 
     __slots__: typing.Sequence[str] = (
         "_close_event",
+        "_entity_factory",
         "_executor",
         "_http_settings",
         "_is_closing",
+        "_on_shutdown",
+        "_on_startup",
         "_proxy_settings",
-        "_entity_factory",
         "_rest",
         "_server",
     )
@@ -262,6 +264,8 @@ def __init__(
         self._executor = executor
         self._http_settings = http_settings if http_settings is not None else config_impl.HTTPSettings()
         self._is_closing = False
+        self._on_shutdown: typing.List[typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]]] = []
+        self._on_startup: typing.List[typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]]] = []
         self._proxy_settings = proxy_settings if proxy_settings is not None else config_impl.ProxySettings()
 
         # Entity creation
@@ -296,6 +300,16 @@ def is_alive(self) -> bool:
     def interaction_server(self) -> interaction_server_.InteractionServer:
         return self._server
 
+    @property
+    def on_shutdown(
+        self,
+    ) -> typing.Sequence[typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]]]:
+        return self._on_shutdown.copy()
+
+    @property
+    def on_startup(self) -> typing.Sequence[typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]]]:
+        return self._on_startup.copy()
+
     @property
     def rest(self) -> rest_api.RESTClient:
         return self._rest
@@ -358,6 +372,26 @@ def print_banner(
         """
         ux.print_banner(banner, allow_color, force_color, extra_args=extra_args)
 
+    def add_shutdown_callback(
+        self, callback: typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        self._on_shutdown.append(callback)
+
+    def remove_shutdown_callback(
+        self, callback: typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        self._on_shutdown.remove(callback)
+
+    def add_startup_callback(
+        self, callback: typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        self._on_startup.append(callback)
+
+    def remove_startup_callback(
+        self, callback: typing.Callable[[RESTBot], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        self._on_startup.remove(callback)
+
     async def close(self) -> None:
         if not self._close_event:
             raise errors.ComponentStateConflictError("Cannot close an inactive bot")
@@ -370,12 +404,18 @@ async def close(self) -> None:
 
         self._is_closing = True
         await self._server.close()
-        await self._rest.close()
-        self._close_event.set()
-        self._close_event = None
-        self._is_closing = False
 
-        _LOGGER.info("bot shut down successfully")
+        try:
+            for callback in self._on_shutdown:
+                await callback(self)
+
+        finally:
+            await self._rest.close()
+            self._close_event.set()
+            self._close_event = None
+            self._is_closing = False
+
+            _LOGGER.info("bot shut down")
 
     async def join(self) -> None:
         if not self._close_event:
@@ -584,6 +624,14 @@ async def start(
             )
 
         self._rest.start()
+        try:
+            for callback in self._on_startup:
+                await callback(self)
+
+        except Exception:
+            await self._rest.close()
+            raise
+
         await self._server.start(
             backlog=backlog,
             host=host,
diff --git a/hikari/traits.py b/hikari/traits.py
index 18040dae..f1bd33a2 100644
--- a/hikari/traits.py
+++ b/hikari/traits.py
@@ -65,6 +65,8 @@
     from hikari.api import shard as gateway_shard
     from hikari.api import voice as voice_
 
+    _RESTBotAwareT = typing.TypeVar("_RESTBotAwareT", bound="RESTBotAware")
+
 
 @typing.runtime_checkable
 class NetworkSettingsAware(fast_protocol.FastProtocolChecking, typing.Protocol):
@@ -566,3 +568,75 @@ class RESTBotAware(InteractionServerAware, Runnable, fast_protocol.FastProtocolC
     """Structural supertype for a component that has all the RESTful components."""
 
     __slots__: typing.Sequence[str] = ()
+
+    @property
+    def on_shutdown(
+        self: _RESTBotAwareT,
+    ) -> typing.Sequence[typing.Callable[[_RESTBotAwareT], typing.Coroutine[typing.Any, typing.Any, None]]]:
+        """Sequence of the bot's asynchronous shutdown callbacks."""
+        raise NotImplementedError
+
+    @property
+    def on_startup(
+        self: _RESTBotAwareT,
+    ) -> typing.Sequence[typing.Callable[[_RESTBotAwareT], typing.Coroutine[typing.Any, typing.Any, None]]]:
+        """Sequence of the bot's asynchronous startup callbacks."""
+        raise NotImplementedError
+
+    def add_shutdown_callback(
+        self, callback: typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        """Add an asynchronous callback to be called when the bot shuts down.
+
+        Parameters
+        ----------
+        callback : typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]]
+            The asynchronous shutdown callback to add.
+        """
+        raise NotImplementedError
+
+    def remove_shutdown_callback(
+        self, callback: typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        """Remove an asynchronous shutdown callback from the bot.
+
+        Parameters
+        ----------
+        callback : typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]]
+            The shutdown callback to remove.
+
+        Raises
+        ------
+        ValueError
+            If the callback was not registered.
+        """
+        raise NotImplementedError
+
+    def add_startup_callback(
+        self, callback: typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        """Add an asynchronous callback to be called when the bot starts up.
+
+        Parameters
+        ----------
+        callback : typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]]
+            The asynchronous startup callback to add.
+        """
+        raise NotImplementedError
+
+    def remove_startup_callback(
+        self, callback: typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]], /
+    ) -> None:
+        """Remove an asynchronous startup callback from the bot.
+
+        Parameters
+        ----------
+        callback : typing.Callable[[RESTBotAware], typing.Coroutine[typing.Any, typing.Any, None]]
+            The asynchronous startup callback to remove.
+
+        Raises
+        ------
+        ValueError
+            If the callback was not registered.
+        """
+        raise NotImplementedError


Test patch:
diff --git a/tests/hikari/impl/test_rest_bot.py b/tests/hikari/impl/test_rest_bot.py
index 3ac09533..880696e7 100644
--- a/tests/hikari/impl/test_rest_bot.py
+++ b/tests/hikari/impl/test_rest_bot.py
@@ -237,11 +237,57 @@ def test_print_banner(self, mock_rest_bot):
 
             print_banner.assert_called_once_with("okokok", True, False, extra_args={"test_key": "test_value"})
 
+    def test_add_shutdown_callback(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+        mock_rest_bot.add_shutdown_callback(callback)
+
+        assert callback in mock_rest_bot.on_shutdown
+
+    def test_remove_shutdown_callback(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+        mock_rest_bot.add_shutdown_callback(callback)
+
+        mock_rest_bot.remove_shutdown_callback(callback)
+
+        assert callback not in mock_rest_bot.on_shutdown
+
+    def test_remove_shutdown_callback_when_not_present(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+
+        with pytest.raises(ValueError, match=".*"):
+            mock_rest_bot.remove_shutdown_callback(callback)
+
+    def test_add_startup_callback(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+        mock_rest_bot.add_startup_callback(callback)
+
+        assert callback in mock_rest_bot.on_startup
+
+    def test_remove_startup_callback(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+        mock_rest_bot.add_startup_callback(callback)
+
+        mock_rest_bot.remove_startup_callback(callback)
+
+        assert callback not in mock_rest_bot.on_startup
+
+    def test_remove_startup_callback_when_not_present(self, mock_rest_bot: rest_bot_impl.RESTBot):
+        callback = mock.Mock()
+
+        with pytest.raises(ValueError, match=".*"):
+            mock_rest_bot.remove_startup_callback(callback)
+
     @pytest.mark.asyncio()
-    async def test_close(self, mock_rest_bot, mock_interaction_server, mock_rest_client):
+    async def test_close(
+        self, mock_rest_bot: rest_bot_impl.RESTBot, mock_interaction_server: mock.Mock, mock_rest_client: mock.Mock
+    ):
+        mock_shutdown_1 = mock.AsyncMock()
+        mock_shutdown_2 = mock.AsyncMock()
         mock_rest_bot._close_event = close_event = mock.Mock()
         mock_interaction_server.close = mock.AsyncMock()
         mock_rest_bot._is_closing = False
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_1)
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_2)
 
         await mock_rest_bot.close()
 
@@ -249,13 +295,45 @@ async def test_close(self, mock_rest_bot, mock_interaction_server, mock_rest_cli
         mock_rest_client.close.assert_awaited_once()
         close_event.set.assert_called_once()
         assert mock_rest_bot._is_closing is False
+        mock_shutdown_1.assert_awaited_once_with(mock_rest_bot)
+        mock_shutdown_2.assert_awaited_once_with(mock_rest_bot)
 
     @pytest.mark.asyncio()
-    async def test_close_when_is_closing(self, mock_rest_bot, mock_interaction_server, mock_rest_client):
+    async def test_close_when_shutdown_callback_raises(
+        self, mock_rest_bot: rest_bot_impl.RESTBot, mock_interaction_server: mock.Mock, mock_rest_client: mock.Mock
+    ):
+        mock_error = KeyError("Too many catgirls")
+        mock_shutdown_1 = mock.AsyncMock(side_effect=mock_error)
+        mock_shutdown_2 = mock.AsyncMock()
+        mock_rest_bot._close_event = close_event = mock.Mock()
+        mock_interaction_server.close = mock.AsyncMock()
+        mock_rest_bot._is_closing = False
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_1)
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_2)
+
+        with pytest.raises(KeyError) as exc_info:
+            await mock_rest_bot.close()
+
+        assert exc_info.value is mock_error
+        mock_interaction_server.close.assert_awaited_once()
+        mock_rest_client.close.assert_awaited_once()
+        close_event.set.assert_called_once()
+        assert mock_rest_bot._is_closing is False
+        mock_shutdown_1.assert_awaited_once_with(mock_rest_bot)
+        mock_shutdown_2.assert_not_called()
+
+    @pytest.mark.asyncio()
+    async def test_close_when_is_closing(
+        self, mock_rest_bot: rest_bot_impl.RESTBot, mock_interaction_server: mock.Mock, mock_rest_client: mock.Mock
+    ):
+        mock_shutdown_1 = mock.AsyncMock()
+        mock_shutdown_2 = mock.AsyncMock()
         mock_rest_bot._close_event = mock.Mock()
         mock_interaction_server.close = mock.AsyncMock()
         mock_rest_bot._is_closing = True
         mock_rest_bot.join = mock.AsyncMock()
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_1)
+        mock_rest_bot.add_shutdown_callback(mock_shutdown_2)
 
         await mock_rest_bot.close()
 
@@ -264,6 +342,8 @@ async def test_close_when_is_closing(self, mock_rest_bot, mock_interaction_serve
         mock_rest_bot._close_event.set.assert_not_called()
         mock_rest_bot.join.assert_awaited_once()
         assert mock_rest_bot._is_closing is True
+        mock_shutdown_1.assert_not_called()
+        mock_shutdown_2.assert_not_called()
 
     @pytest.mark.asyncio()
     async def test_close_when_inactive(self, mock_rest_bot):
@@ -440,9 +520,15 @@ def test_run_ignores_close_executor_when_not_present(self, mock_rest_bot):
         assert mock_rest_bot.executor is None
 
     @pytest.mark.asyncio()
-    async def test_start(self, mock_rest_bot, mock_interaction_server, mock_rest_client):
+    async def test_start(
+        self, mock_rest_bot: rest_bot_impl.RESTBot, mock_interaction_server: mock.Mock, mock_rest_client: mock.Mock
+    ):
         mock_socket = object()
         mock_ssl_context = object()
+        mock_callback_1 = mock.AsyncMock()
+        mock_callback_2 = mock.AsyncMock()
+        mock_rest_bot.add_startup_callback(mock_callback_1)
+        mock_rest_bot.add_startup_callback(mock_callback_2)
         mock_rest_bot._is_closing = True
 
         with mock.patch.object(ux, "check_for_updates"):
@@ -473,7 +559,48 @@ async def test_start(self, mock_rest_bot, mock_interaction_server, mock_rest_cli
             ssl_context=mock_ssl_context,
         )
         mock_rest_client.start.assert_called_once_with()
+        mock_rest_client.close.assert_not_called()
+        assert mock_rest_bot._is_closing is False
+        mock_callback_1.assert_awaited_once_with(mock_rest_bot)
+        mock_callback_2.assert_awaited_once_with(mock_rest_bot)
+
+    @pytest.mark.asyncio()
+    async def test_start_when_startup_callback_raises(
+        self, mock_rest_bot: rest_bot_impl.RESTBot, mock_interaction_server: mock.Mock, mock_rest_client: mock.Mock
+    ):
+        mock_socket = object()
+        mock_ssl_context = object()
+        mock_rest_bot._is_closing = True
+        mock_error = TypeError("Not a real catgirl")
+        mock_callback_1 = mock.AsyncMock(side_effect=mock_error)
+        mock_callback_2 = mock.AsyncMock()
+        mock_rest_bot.add_startup_callback(mock_callback_1)
+        mock_rest_bot.add_startup_callback(mock_callback_2)
+
+        with mock.patch.object(ux, "check_for_updates"):
+            with pytest.raises(TypeError) as exc_info:
+                await mock_rest_bot.start(
+                    backlog=34123,
+                    check_for_updates=False,
+                    host="hostostosot",
+                    port=123123123,
+                    path="patpatpapt",
+                    reuse_address=True,
+                    reuse_port=False,
+                    socket=mock_socket,
+                    shutdown_timeout=4312312.3132132,
+                    ssl_context=mock_ssl_context,
+                )
+
+            assert exc_info.value is mock_error
+            ux.check_for_updates.assert_not_called()
+
+        mock_interaction_server.start.assert_not_called()
+        mock_rest_client.start.assert_called_once_with()
+        mock_rest_client.close.assert_awaited_once_with()
         assert mock_rest_bot._is_closing is False
+        mock_callback_1.assert_awaited_once_with(mock_rest_bot)
+        mock_callback_2.assert_not_called()
 
     @pytest.mark.asyncio()
     async def test_start_checks_for_update(self, mock_rest_bot, mock_http_settings, mock_proxy_settings):