test_mqtt.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276
  1. # systemctl-mqtt - MQTT client triggering shutdown on systemd-based systems
  2. #
  3. # Copyright (C) 2020 Fabian Peter Hammerle <fabian@hammerle.me>
  4. #
  5. # This program is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU General Public License as published by
  7. # the Free Software Foundation, either version 3 of the License, or
  8. # any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU General Public License
  16. # along with this program. If not, see <https://www.gnu.org/licenses/>.
  17. import logging
  18. import threading
  19. import time
  20. import unittest.mock
  21. import paho.mqtt.client
  22. import pytest
  23. from paho.mqtt.client import MQTTMessage
  24. import systemctl_mqtt
  25. # pylint: disable=protected-access
  26. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  27. @pytest.mark.parametrize("mqtt_port", [1833])
  28. @pytest.mark.parametrize("mqtt_topic_prefix", ["systemctl/host", "system/command"])
  29. def test__run(caplog, mqtt_host, mqtt_port, mqtt_topic_prefix):
  30. caplog.set_level(logging.DEBUG)
  31. with unittest.mock.patch(
  32. "socket.create_connection"
  33. ) as create_socket_mock, unittest.mock.patch(
  34. "ssl.SSLContext.wrap_socket", autospec=True,
  35. ) as ssl_wrap_socket_mock, unittest.mock.patch(
  36. "paho.mqtt.client.Client.loop_forever", autospec=True,
  37. ) as mqtt_loop_forever_mock, unittest.mock.patch(
  38. "gi.repository.GLib.MainLoop.run"
  39. ) as glib_loop_mock, unittest.mock.patch(
  40. "systemctl_mqtt._get_login_manager"
  41. ):
  42. ssl_wrap_socket_mock.return_value.send = len
  43. systemctl_mqtt._run(
  44. mqtt_host=mqtt_host,
  45. mqtt_port=mqtt_port,
  46. mqtt_username=None,
  47. mqtt_password=None,
  48. mqtt_topic_prefix=mqtt_topic_prefix,
  49. )
  50. assert caplog.records[0].levelno == logging.INFO
  51. assert caplog.records[0].message == "connecting to MQTT broker {}:{}".format(
  52. mqtt_host, mqtt_port
  53. )
  54. # correct remote?
  55. assert create_socket_mock.call_count == 1
  56. create_socket_args, _ = create_socket_mock.call_args
  57. assert create_socket_args[0] == (mqtt_host, mqtt_port)
  58. # ssl enabled?
  59. assert ssl_wrap_socket_mock.call_count == 1
  60. ssl_context = ssl_wrap_socket_mock.call_args[0][0] # self
  61. assert ssl_context.check_hostname is True
  62. assert ssl_wrap_socket_mock.call_args[1]["server_hostname"] == mqtt_host
  63. # loop started?
  64. while threading.active_count() > 1:
  65. time.sleep(0.01)
  66. assert mqtt_loop_forever_mock.call_count == 1
  67. (mqtt_client,) = mqtt_loop_forever_mock.call_args[0]
  68. assert mqtt_client._tls_insecure is False
  69. # credentials
  70. assert mqtt_client._username is None
  71. assert mqtt_client._password is None
  72. # connect callback
  73. caplog.clear()
  74. mqtt_client.socket().getpeername.return_value = (mqtt_host, mqtt_port)
  75. with unittest.mock.patch(
  76. "paho.mqtt.client.Client.subscribe"
  77. ) as mqtt_subscribe_mock:
  78. mqtt_client.on_connect(mqtt_client, mqtt_client._userdata, {}, 0)
  79. state = mqtt_client._userdata
  80. assert (
  81. state._login_manager.connect_to_signal.call_args[1]["signal_name"]
  82. == "PrepareForShutdown"
  83. )
  84. mqtt_subscribe_mock.assert_called_once_with(mqtt_topic_prefix + "/poweroff")
  85. assert mqtt_client.on_message is None
  86. assert ( # pylint: disable=comparison-with-callable
  87. mqtt_client._on_message_filtered[mqtt_topic_prefix + "/poweroff"]
  88. == systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  89. "poweroff"
  90. ].mqtt_message_callback
  91. )
  92. assert caplog.records[0].levelno == logging.DEBUG
  93. assert caplog.records[0].message == "connected to MQTT broker {}:{}".format(
  94. mqtt_host, mqtt_port
  95. )
  96. assert caplog.records[1].levelno == logging.DEBUG
  97. assert caplog.records[1].message == "acquired shutdown inhibitor lock"
  98. assert caplog.records[2].levelno == logging.INFO
  99. assert caplog.records[2].message == "subscribing to {}".format(
  100. mqtt_topic_prefix + "/poweroff"
  101. )
  102. assert caplog.records[3].levelno == logging.DEBUG
  103. assert caplog.records[3].message == "registered MQTT callback for topic {}".format(
  104. mqtt_topic_prefix + "/poweroff"
  105. ) + " triggering {}".format(
  106. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"].action
  107. )
  108. # dbus loop started?
  109. glib_loop_mock.assert_called_once_with()
  110. # waited for mqtt loop to stop?
  111. assert mqtt_client._thread_terminate
  112. assert mqtt_client._thread is None
  113. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  114. @pytest.mark.parametrize("mqtt_port", [1833])
  115. @pytest.mark.parametrize("mqtt_username", ["me"])
  116. @pytest.mark.parametrize("mqtt_password", [None, "secret"])
  117. @pytest.mark.parametrize("mqtt_topic_prefix", ["systemctl/host"])
  118. def test__run_authentication(
  119. mqtt_host, mqtt_port, mqtt_username, mqtt_password, mqtt_topic_prefix
  120. ):
  121. with unittest.mock.patch("socket.create_connection"), unittest.mock.patch(
  122. "ssl.SSLContext.wrap_socket"
  123. ) as ssl_wrap_socket_mock, unittest.mock.patch(
  124. "paho.mqtt.client.Client.loop_forever", autospec=True,
  125. ) as mqtt_loop_forever_mock, unittest.mock.patch(
  126. "gi.repository.GLib.MainLoop.run"
  127. ), unittest.mock.patch(
  128. "systemctl_mqtt._get_login_manager"
  129. ):
  130. ssl_wrap_socket_mock.return_value.send = len
  131. systemctl_mqtt._run(
  132. mqtt_host=mqtt_host,
  133. mqtt_port=mqtt_port,
  134. mqtt_username=mqtt_username,
  135. mqtt_password=mqtt_password,
  136. mqtt_topic_prefix=mqtt_topic_prefix,
  137. )
  138. assert mqtt_loop_forever_mock.call_count == 1
  139. (mqtt_client,) = mqtt_loop_forever_mock.call_args[0]
  140. assert mqtt_client._username.decode() == mqtt_username
  141. if mqtt_password:
  142. assert mqtt_client._password.decode() == mqtt_password
  143. else:
  144. assert mqtt_client._password is None
  145. def _initialize_mqtt_client(
  146. mqtt_host, mqtt_port, mqtt_topic_prefix
  147. ) -> paho.mqtt.client.Client:
  148. with unittest.mock.patch("socket.create_connection"), unittest.mock.patch(
  149. "ssl.SSLContext.wrap_socket",
  150. ) as ssl_wrap_socket_mock, unittest.mock.patch(
  151. "paho.mqtt.client.Client.loop_forever", autospec=True,
  152. ) as mqtt_loop_forever_mock, unittest.mock.patch(
  153. "gi.repository.GLib.MainLoop.run"
  154. ), unittest.mock.patch(
  155. "systemctl_mqtt._get_login_manager"
  156. ):
  157. ssl_wrap_socket_mock.return_value.send = len
  158. systemctl_mqtt._run(
  159. mqtt_host=mqtt_host,
  160. mqtt_port=mqtt_port,
  161. mqtt_username=None,
  162. mqtt_password=None,
  163. mqtt_topic_prefix=mqtt_topic_prefix,
  164. )
  165. while threading.active_count() > 1:
  166. time.sleep(0.01)
  167. assert mqtt_loop_forever_mock.call_count == 1
  168. (mqtt_client,) = mqtt_loop_forever_mock.call_args[0]
  169. mqtt_client.socket().getpeername.return_value = (mqtt_host, mqtt_port)
  170. mqtt_client.on_connect(mqtt_client, mqtt_client._userdata, {}, 0)
  171. return mqtt_client
  172. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  173. @pytest.mark.parametrize("mqtt_port", [1833])
  174. @pytest.mark.parametrize("mqtt_topic_prefix", ["systemctl/host", "system/command"])
  175. def test__client_handle_message(caplog, mqtt_host, mqtt_port, mqtt_topic_prefix):
  176. mqtt_client = _initialize_mqtt_client(
  177. mqtt_host=mqtt_host, mqtt_port=mqtt_port, mqtt_topic_prefix=mqtt_topic_prefix
  178. )
  179. caplog.clear()
  180. caplog.set_level(logging.DEBUG)
  181. poweroff_message = MQTTMessage(topic=mqtt_topic_prefix.encode() + b"/poweroff")
  182. with unittest.mock.patch.object(
  183. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  184. ) as poweroff_action_mock:
  185. mqtt_client._handle_on_message(poweroff_message)
  186. poweroff_action_mock.assert_called_once_with()
  187. assert all(r.levelno == logging.DEBUG for r in caplog.records)
  188. assert caplog.records[0].message == "received topic={} payload=b''".format(
  189. poweroff_message.topic
  190. )
  191. assert caplog.records[1].message.startswith("executing action poweroff")
  192. assert caplog.records[2].message.startswith("completed action poweroff")
  193. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  194. @pytest.mark.parametrize("mqtt_port", [1833])
  195. @pytest.mark.parametrize("mqtt_password", ["secret"])
  196. def test__run_authentication_missing_username(mqtt_host, mqtt_port, mqtt_password):
  197. with unittest.mock.patch("paho.mqtt.client.Client"), unittest.mock.patch(
  198. "systemctl_mqtt._get_login_manager"
  199. ):
  200. with pytest.raises(ValueError):
  201. systemctl_mqtt._run(
  202. mqtt_host=mqtt_host,
  203. mqtt_port=mqtt_port,
  204. mqtt_username=None,
  205. mqtt_password=mqtt_password,
  206. mqtt_topic_prefix="prefix",
  207. )
  208. @pytest.mark.parametrize("mqtt_topic", ["system/command/poweroff"])
  209. @pytest.mark.parametrize("payload", [b"", b"junk"])
  210. def test_mqtt_message_callback_poweroff(caplog, mqtt_topic: str, payload: bytes):
  211. message = MQTTMessage(topic=mqtt_topic.encode())
  212. message.payload = payload
  213. with unittest.mock.patch.object(
  214. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  215. ) as action_mock, caplog.at_level(logging.DEBUG):
  216. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  217. "poweroff"
  218. ].mqtt_message_callback(
  219. None, None, message # type: ignore
  220. )
  221. action_mock.assert_called_once_with()
  222. assert len(caplog.records) == 3
  223. assert caplog.records[0].levelno == logging.DEBUG
  224. assert caplog.records[0].message == (
  225. "received topic={} payload={!r}".format(mqtt_topic, payload)
  226. )
  227. assert caplog.records[1].levelno == logging.DEBUG
  228. assert caplog.records[1].message.startswith(
  229. "executing action {} ({!r})".format("poweroff", action_mock)
  230. )
  231. assert caplog.records[2].levelno == logging.DEBUG
  232. assert caplog.records[2].message.startswith(
  233. "completed action {} ({!r})".format("poweroff", action_mock)
  234. )
  235. @pytest.mark.parametrize("mqtt_topic", ["system/command/poweroff"])
  236. @pytest.mark.parametrize("payload", [b"", b"junk"])
  237. def test_mqtt_message_callback_poweroff_retained(
  238. caplog, mqtt_topic: str, payload: bytes
  239. ):
  240. message = MQTTMessage(topic=mqtt_topic.encode())
  241. message.payload = payload
  242. message.retain = True
  243. with unittest.mock.patch.object(
  244. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  245. ) as action_mock, caplog.at_level(logging.DEBUG):
  246. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  247. "poweroff"
  248. ].mqtt_message_callback(
  249. None, None, message # type: ignore
  250. )
  251. action_mock.assert_not_called()
  252. assert len(caplog.records) == 2
  253. assert caplog.records[0].levelno == logging.DEBUG
  254. assert caplog.records[0].message == (
  255. "received topic={} payload={!r}".format(mqtt_topic, payload)
  256. )
  257. assert caplog.records[1].levelno == logging.INFO
  258. assert caplog.records[1].message == "ignoring retained message"