subscribe to logind's PrepareForShutdown signal

systemctl_mqtt/__init__.py

@@ -115,6 +115,11 @@ def _schedule_shutdown(action: str) -> None:
 class _State:
     def __init__(self, mqtt_topic_prefix: str) -> None:
         self._mqtt_topic_prefix = mqtt_topic_prefix
+        self._login_manager = _get_login_manager()  # type: dbus.proxies.Interface
+        self._login_manager.connect_to_signal(
+            signal_name="PrepareForShutdown",
+            handler_function=self.prepare_for_shutdown_handler,
+        )
         self._shutdown_lock = None  # type: typing.Optional[dbus.types.UnixFd]
         self._shutdown_lock_mutex = threading.Lock()
 
@@ -139,6 +144,14 @@ class _State:
                 _LOGGER.debug("released shutdown inhibitor lock")
                 self._shutdown_lock = None
 
+    def prepare_for_shutdown_handler(self, active: bool) -> None:
+        if active:
+            _LOGGER.debug("system preparing for shutdown")
+            self.release_shutdown_lock()
+        else:
+            _LOGGER.debug("system shutdown failed?")
+            self.acquire_shutdown_lock()
+
 
 class _MQTTAction:
 
@@ -204,6 +217,8 @@ def _run(
     mqtt_password: typing.Optional[str],
     mqtt_topic_prefix: str,
 ) -> None:
+    # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
+    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
     # https://pypi.org/project/paho-mqtt/
     mqtt_client = paho.mqtt.client.Client(
         userdata=_State(mqtt_topic_prefix=mqtt_topic_prefix)
@@ -223,8 +238,6 @@ def _run(
     # loop_forever attempts to reconnect if disconnected
     # https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1744
     mqtt_client.loop_start()
-    # https://dbus.freedesktop.org/doc/dbus-python/tutorial.html#setting-up-an-event-loop
-    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
     try:
         gi.repository.GLib.MainLoop().run()
     finally:

tests/test_mqtt.py

@@ -42,7 +42,9 @@ def test__run(caplog, mqtt_host, mqtt_port, mqtt_topic_prefix):
         "paho.mqtt.client.Client.loop_forever", autospec=True,
     ) as mqtt_loop_forever_mock, unittest.mock.patch(
         "gi.repository.GLib.MainLoop.run"
-    ) as glib_loop_mock:
+    ) as glib_loop_mock, unittest.mock.patch(
+        "systemctl_mqtt._get_login_manager"
+    ):
         ssl_wrap_socket_mock.return_value.send = len
         systemctl_mqtt._run(
             mqtt_host=mqtt_host,
@@ -126,6 +128,8 @@ def test__run_authentication(
         "paho.mqtt.client.Client.loop_forever", autospec=True,
     ) as mqtt_loop_forever_mock, unittest.mock.patch(
         "gi.repository.GLib.MainLoop.run"
+    ), unittest.mock.patch(
+        "systemctl_mqtt._get_login_manager"
     ):
         ssl_wrap_socket_mock.return_value.send = len
         systemctl_mqtt._run(
@@ -153,6 +157,8 @@ def _initialize_mqtt_client(
         "paho.mqtt.client.Client.loop_forever", autospec=True,
     ) as mqtt_loop_forever_mock, unittest.mock.patch(
         "gi.repository.GLib.MainLoop.run"
+    ), unittest.mock.patch(
+        "systemctl_mqtt._get_login_manager"
     ):
         ssl_wrap_socket_mock.return_value.send = len
         systemctl_mqtt._run(
@@ -198,7 +204,9 @@ def test__client_handle_message(caplog, mqtt_host, mqtt_port, mqtt_topic_prefix)
 @pytest.mark.parametrize("mqtt_port", [1833])
 @pytest.mark.parametrize("mqtt_password", ["secret"])
 def test__run_authentication_missing_username(mqtt_host, mqtt_port, mqtt_password):
-    with unittest.mock.patch("paho.mqtt.client.Client"):
+    with unittest.mock.patch("paho.mqtt.client.Client"), unittest.mock.patch(
+        "systemctl_mqtt._get_login_manager"
+    ):
         with pytest.raises(ValueError):
             systemctl_mqtt._run(
                 mqtt_host=mqtt_host,
@@ -262,22 +270,3 @@ def test_mqtt_message_callback_poweroff_retained(
     )
     assert caplog.records[1].levelno == logging.INFO
     assert caplog.records[1].message == "ignoring retained message"
-
-
-def test_shutdown_lock():
-    state = systemctl_mqtt._State(mqtt_topic_prefix="any")
-    lock_fd = unittest.mock.MagicMock()
-    with unittest.mock.patch(
-        "systemctl_mqtt._get_login_manager"
-    ) as get_login_manager_mock:
-        get_login_manager_mock.return_value.Inhibit.return_value = lock_fd
-        state.acquire_shutdown_lock()
-    get_login_manager_mock.return_value.Inhibit.assert_called_once_with(
-        "shutdown", "systemctl-mqtt", "Report shutdown via MQTT", "delay",
-    )
-    assert state._shutdown_lock == lock_fd
-    # https://dbus.freedesktop.org/doc/dbus-python/dbus.types.html#dbus.types.UnixFd.take
-    lock_fd.take.return_value = "fdnum"
-    with unittest.mock.patch("os.close") as close_mock:
-        state.release_shutdown_lock()
-    close_mock.assert_called_once_with("fdnum")

tests/test_state_dbus.py

@@ -0,0 +1,64 @@
+# systemctl-mqtt - MQTT client triggering shutdown on systemd-based systems
+#
+# Copyright (C) 2020 Fabian Peter Hammerle <fabian@hammerle.me>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import unittest.mock
+
+import pytest
+
+import systemctl_mqtt
+
+# pylint: disable=protected-access
+
+
+def test_shutdown_lock():
+    lock_fd = unittest.mock.MagicMock()
+    with unittest.mock.patch("systemctl_mqtt._get_login_manager"):
+        state = systemctl_mqtt._State(mqtt_topic_prefix="any")
+        state._login_manager.Inhibit.return_value = lock_fd
+        state.acquire_shutdown_lock()
+    state._login_manager.Inhibit.assert_called_once_with(
+        "shutdown", "systemctl-mqtt", "Report shutdown via MQTT", "delay",
+    )
+    assert state._shutdown_lock == lock_fd
+    # https://dbus.freedesktop.org/doc/dbus-python/dbus.types.html#dbus.types.UnixFd.take
+    lock_fd.take.return_value = "fdnum"
+    with unittest.mock.patch("os.close") as close_mock:
+        state.release_shutdown_lock()
+    close_mock.assert_called_once_with("fdnum")
+
+
+@pytest.mark.parametrize("active", [True, False])
+def test_prepare_for_shutdown_handler(active):
+    with unittest.mock.patch("systemctl_mqtt._get_login_manager"):
+        state = systemctl_mqtt._State(mqtt_topic_prefix="any")
+    # pylint: disable=no-member,comparison-with-callable
+    connect_to_signal_kwargs = state._login_manager.connect_to_signal.call_args[1]
+    assert connect_to_signal_kwargs["signal_name"] == "PrepareForShutdown"
+    handler_function = connect_to_signal_kwargs["handler_function"]
+    assert handler_function == state.prepare_for_shutdown_handler
+    with unittest.mock.patch.object(
+        state, "acquire_shutdown_lock"
+    ) as acquire_lock_mock, unittest.mock.patch.object(
+        state, "release_shutdown_lock"
+    ) as release_lock_mock:
+        handler_function(active)
+    if active:
+        acquire_lock_mock.assert_not_called()
+        release_lock_mock.assert_called_once_with()
+    else:
+        acquire_lock_mock.assert_called_once_with()
+        release_lock_mock.assert_not_called()