1
0

test_mqtt.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272
  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, unittest.mock.patch.object(
  78. mqtt_client._userdata, "acquire_shutdown_lock"
  79. ) as acquire_shutdown_lock_mock:
  80. mqtt_client.on_connect(mqtt_client, mqtt_client._userdata, {}, 0)
  81. acquire_shutdown_lock_mock.assert_called_once_with()
  82. mqtt_subscribe_mock.assert_called_once_with(mqtt_topic_prefix + "/poweroff")
  83. assert mqtt_client.on_message is None
  84. assert ( # pylint: disable=comparison-with-callable
  85. mqtt_client._on_message_filtered[mqtt_topic_prefix + "/poweroff"]
  86. == systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  87. "poweroff"
  88. ].mqtt_message_callback
  89. )
  90. assert caplog.records[0].levelno == logging.DEBUG
  91. assert caplog.records[0].message == "connected to MQTT broker {}:{}".format(
  92. mqtt_host, mqtt_port
  93. )
  94. assert caplog.records[1].levelno == logging.INFO
  95. assert caplog.records[1].message == "subscribing to {}".format(
  96. mqtt_topic_prefix + "/poweroff"
  97. )
  98. assert caplog.records[2].levelno == logging.DEBUG
  99. assert caplog.records[2].message == "registered MQTT callback for topic {}".format(
  100. mqtt_topic_prefix + "/poweroff"
  101. ) + " triggering {}".format(
  102. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"].action
  103. )
  104. # dbus loop started?
  105. glib_loop_mock.assert_called_once_with()
  106. # waited for mqtt loop to stop?
  107. assert mqtt_client._thread_terminate
  108. assert mqtt_client._thread is None
  109. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  110. @pytest.mark.parametrize("mqtt_port", [1833])
  111. @pytest.mark.parametrize("mqtt_username", ["me"])
  112. @pytest.mark.parametrize("mqtt_password", [None, "secret"])
  113. @pytest.mark.parametrize("mqtt_topic_prefix", ["systemctl/host"])
  114. def test__run_authentication(
  115. mqtt_host, mqtt_port, mqtt_username, mqtt_password, mqtt_topic_prefix
  116. ):
  117. with unittest.mock.patch("socket.create_connection"), unittest.mock.patch(
  118. "ssl.SSLContext.wrap_socket"
  119. ) as ssl_wrap_socket_mock, unittest.mock.patch(
  120. "paho.mqtt.client.Client.loop_forever", autospec=True,
  121. ) as mqtt_loop_forever_mock, unittest.mock.patch(
  122. "gi.repository.GLib.MainLoop.run"
  123. ), unittest.mock.patch(
  124. "systemctl_mqtt._get_login_manager"
  125. ):
  126. ssl_wrap_socket_mock.return_value.send = len
  127. systemctl_mqtt._run(
  128. mqtt_host=mqtt_host,
  129. mqtt_port=mqtt_port,
  130. mqtt_username=mqtt_username,
  131. mqtt_password=mqtt_password,
  132. mqtt_topic_prefix=mqtt_topic_prefix,
  133. )
  134. assert mqtt_loop_forever_mock.call_count == 1
  135. (mqtt_client,) = mqtt_loop_forever_mock.call_args[0]
  136. assert mqtt_client._username.decode() == mqtt_username
  137. if mqtt_password:
  138. assert mqtt_client._password.decode() == mqtt_password
  139. else:
  140. assert mqtt_client._password is None
  141. def _initialize_mqtt_client(
  142. mqtt_host, mqtt_port, mqtt_topic_prefix
  143. ) -> paho.mqtt.client.Client:
  144. with unittest.mock.patch("socket.create_connection"), unittest.mock.patch(
  145. "ssl.SSLContext.wrap_socket",
  146. ) as ssl_wrap_socket_mock, unittest.mock.patch(
  147. "paho.mqtt.client.Client.loop_forever", autospec=True,
  148. ) as mqtt_loop_forever_mock, unittest.mock.patch(
  149. "gi.repository.GLib.MainLoop.run"
  150. ), unittest.mock.patch(
  151. "systemctl_mqtt._get_login_manager"
  152. ):
  153. ssl_wrap_socket_mock.return_value.send = len
  154. systemctl_mqtt._run(
  155. mqtt_host=mqtt_host,
  156. mqtt_port=mqtt_port,
  157. mqtt_username=None,
  158. mqtt_password=None,
  159. mqtt_topic_prefix=mqtt_topic_prefix,
  160. )
  161. while threading.active_count() > 1:
  162. time.sleep(0.01)
  163. assert mqtt_loop_forever_mock.call_count == 1
  164. (mqtt_client,) = mqtt_loop_forever_mock.call_args[0]
  165. mqtt_client.socket().getpeername.return_value = (mqtt_host, mqtt_port)
  166. mqtt_client.on_connect(mqtt_client, mqtt_client._userdata, {}, 0)
  167. return mqtt_client
  168. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  169. @pytest.mark.parametrize("mqtt_port", [1833])
  170. @pytest.mark.parametrize("mqtt_topic_prefix", ["systemctl/host", "system/command"])
  171. def test__client_handle_message(caplog, mqtt_host, mqtt_port, mqtt_topic_prefix):
  172. mqtt_client = _initialize_mqtt_client(
  173. mqtt_host=mqtt_host, mqtt_port=mqtt_port, mqtt_topic_prefix=mqtt_topic_prefix
  174. )
  175. caplog.clear()
  176. caplog.set_level(logging.DEBUG)
  177. poweroff_message = MQTTMessage(topic=mqtt_topic_prefix.encode() + b"/poweroff")
  178. with unittest.mock.patch.object(
  179. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  180. ) as poweroff_action_mock:
  181. mqtt_client._handle_on_message(poweroff_message)
  182. poweroff_action_mock.assert_called_once_with()
  183. assert all(r.levelno == logging.DEBUG for r in caplog.records)
  184. assert caplog.records[0].message == "received topic={} payload=b''".format(
  185. poweroff_message.topic
  186. )
  187. assert caplog.records[1].message.startswith("executing action poweroff")
  188. assert caplog.records[2].message.startswith("completed action poweroff")
  189. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  190. @pytest.mark.parametrize("mqtt_port", [1833])
  191. @pytest.mark.parametrize("mqtt_password", ["secret"])
  192. def test__run_authentication_missing_username(mqtt_host, mqtt_port, mqtt_password):
  193. with unittest.mock.patch("paho.mqtt.client.Client"), unittest.mock.patch(
  194. "systemctl_mqtt._get_login_manager"
  195. ):
  196. with pytest.raises(ValueError):
  197. systemctl_mqtt._run(
  198. mqtt_host=mqtt_host,
  199. mqtt_port=mqtt_port,
  200. mqtt_username=None,
  201. mqtt_password=mqtt_password,
  202. mqtt_topic_prefix="prefix",
  203. )
  204. @pytest.mark.parametrize("mqtt_topic", ["system/command/poweroff"])
  205. @pytest.mark.parametrize("payload", [b"", b"junk"])
  206. def test_mqtt_message_callback_poweroff(caplog, mqtt_topic: str, payload: bytes):
  207. message = MQTTMessage(topic=mqtt_topic.encode())
  208. message.payload = payload
  209. with unittest.mock.patch.object(
  210. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  211. ) as action_mock, caplog.at_level(logging.DEBUG):
  212. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  213. "poweroff"
  214. ].mqtt_message_callback(
  215. None, None, message # type: ignore
  216. )
  217. action_mock.assert_called_once_with()
  218. assert len(caplog.records) == 3
  219. assert caplog.records[0].levelno == logging.DEBUG
  220. assert caplog.records[0].message == (
  221. "received topic={} payload={!r}".format(mqtt_topic, payload)
  222. )
  223. assert caplog.records[1].levelno == logging.DEBUG
  224. assert caplog.records[1].message.startswith(
  225. "executing action {} ({!r})".format("poweroff", action_mock)
  226. )
  227. assert caplog.records[2].levelno == logging.DEBUG
  228. assert caplog.records[2].message.startswith(
  229. "completed action {} ({!r})".format("poweroff", action_mock)
  230. )
  231. @pytest.mark.parametrize("mqtt_topic", ["system/command/poweroff"])
  232. @pytest.mark.parametrize("payload", [b"", b"junk"])
  233. def test_mqtt_message_callback_poweroff_retained(
  234. caplog, mqtt_topic: str, payload: bytes
  235. ):
  236. message = MQTTMessage(topic=mqtt_topic.encode())
  237. message.payload = payload
  238. message.retain = True
  239. with unittest.mock.patch.object(
  240. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING["poweroff"], "action",
  241. ) as action_mock, caplog.at_level(logging.DEBUG):
  242. systemctl_mqtt._MQTT_TOPIC_SUFFIX_ACTION_MAPPING[
  243. "poweroff"
  244. ].mqtt_message_callback(
  245. None, None, message # type: ignore
  246. )
  247. action_mock.assert_not_called()
  248. assert len(caplog.records) == 2
  249. assert caplog.records[0].levelno == logging.DEBUG
  250. assert caplog.records[0].message == (
  251. "received topic={} payload={!r}".format(mqtt_topic, payload)
  252. )
  253. assert caplog.records[1].levelno == logging.INFO
  254. assert caplog.records[1].message == "ignoring retained message"