1
0

test_mqtt.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491
  1. # switchbot-mqtt - MQTT client controlling SwitchBot button & curtain automators,
  2. # compatible with home-assistant.io's MQTT Switch & Cover platform
  3. #
  4. # Copyright (C) 2020 Fabian Peter Hammerle <fabian@hammerle.me>
  5. #
  6. # This program is free software: you can redistribute it and/or modify
  7. # it under the terms of the GNU General Public License as published by
  8. # the Free Software Foundation, either version 3 of the License, or
  9. # any later version.
  10. #
  11. # This program is distributed in the hope that it will be useful,
  12. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. # GNU General Public License for more details.
  15. #
  16. # You should have received a copy of the GNU General Public License
  17. # along with this program. If not, see <https://www.gnu.org/licenses/>.
  18. import logging
  19. import typing
  20. import unittest.mock
  21. import pytest
  22. from paho.mqtt.client import MQTT_ERR_QUEUE_SIZE, MQTT_ERR_SUCCESS, MQTTMessage, Client
  23. import switchbot_mqtt
  24. # pylint: disable=protected-access
  25. # pylint: disable=too-many-arguments; these are tests, no API
  26. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  27. @pytest.mark.parametrize("mqtt_port", [1833])
  28. @pytest.mark.parametrize("retry_count", [3, 21])
  29. @pytest.mark.parametrize(
  30. "device_passwords",
  31. [{}, {"11:22:33:44:55:66": "password", "aa:bb:cc:dd:ee:ff": "secret"}],
  32. )
  33. @pytest.mark.parametrize("fetch_device_info", [True, False])
  34. def test__run(
  35. caplog, mqtt_host, mqtt_port, retry_count, device_passwords, fetch_device_info
  36. ):
  37. with unittest.mock.patch(
  38. "paho.mqtt.client.Client"
  39. ) as mqtt_client_mock, caplog.at_level(logging.DEBUG):
  40. switchbot_mqtt._run(
  41. mqtt_host=mqtt_host,
  42. mqtt_port=mqtt_port,
  43. mqtt_username=None,
  44. mqtt_password=None,
  45. retry_count=retry_count,
  46. device_passwords=device_passwords,
  47. fetch_device_info=fetch_device_info,
  48. )
  49. mqtt_client_mock.assert_called_once_with(
  50. userdata=switchbot_mqtt._MQTTCallbackUserdata(
  51. retry_count=retry_count,
  52. device_passwords=device_passwords,
  53. fetch_device_info=fetch_device_info,
  54. )
  55. )
  56. assert not mqtt_client_mock().username_pw_set.called
  57. mqtt_client_mock().connect.assert_called_once_with(host=mqtt_host, port=mqtt_port)
  58. mqtt_client_mock().socket().getpeername.return_value = (mqtt_host, mqtt_port)
  59. with caplog.at_level(logging.DEBUG):
  60. mqtt_client_mock().on_connect(mqtt_client_mock(), None, {}, 0)
  61. assert mqtt_client_mock().subscribe.call_args_list == [
  62. unittest.mock.call("homeassistant/switch/switchbot/+/set"),
  63. unittest.mock.call("homeassistant/cover/switchbot-curtain/+/set"),
  64. ]
  65. assert mqtt_client_mock().message_callback_add.call_args_list == [
  66. unittest.mock.call(
  67. sub="homeassistant/switch/switchbot/+/set",
  68. callback=switchbot_mqtt._ButtonAutomator._mqtt_command_callback,
  69. ),
  70. unittest.mock.call(
  71. sub="homeassistant/cover/switchbot-curtain/+/set",
  72. callback=switchbot_mqtt._CurtainMotor._mqtt_command_callback,
  73. ),
  74. ]
  75. mqtt_client_mock().loop_forever.assert_called_once_with()
  76. assert caplog.record_tuples == [
  77. (
  78. "switchbot_mqtt",
  79. logging.INFO,
  80. "connecting to MQTT broker {}:{}".format(mqtt_host, mqtt_port),
  81. ),
  82. (
  83. "switchbot_mqtt",
  84. logging.DEBUG,
  85. "connected to MQTT broker {}:{}".format(mqtt_host, mqtt_port),
  86. ),
  87. (
  88. "switchbot_mqtt",
  89. logging.INFO,
  90. "subscribing to MQTT topic 'homeassistant/switch/switchbot/+/set'",
  91. ),
  92. (
  93. "switchbot_mqtt",
  94. logging.INFO,
  95. "subscribing to MQTT topic 'homeassistant/cover/switchbot-curtain/+/set'",
  96. ),
  97. ]
  98. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  99. @pytest.mark.parametrize("mqtt_port", [1833])
  100. @pytest.mark.parametrize("mqtt_username", ["me"])
  101. @pytest.mark.parametrize("mqtt_password", [None, "secret"])
  102. def test__run_authentication(mqtt_host, mqtt_port, mqtt_username, mqtt_password):
  103. with unittest.mock.patch("paho.mqtt.client.Client") as mqtt_client_mock:
  104. switchbot_mqtt._run(
  105. mqtt_host=mqtt_host,
  106. mqtt_port=mqtt_port,
  107. mqtt_username=mqtt_username,
  108. mqtt_password=mqtt_password,
  109. retry_count=7,
  110. device_passwords={},
  111. fetch_device_info=True,
  112. )
  113. mqtt_client_mock.assert_called_once_with(
  114. userdata=switchbot_mqtt._MQTTCallbackUserdata(
  115. retry_count=7, device_passwords={}, fetch_device_info=True
  116. )
  117. )
  118. mqtt_client_mock().username_pw_set.assert_called_once_with(
  119. username=mqtt_username, password=mqtt_password
  120. )
  121. @pytest.mark.parametrize("mqtt_host", ["mqtt-broker.local"])
  122. @pytest.mark.parametrize("mqtt_port", [1833])
  123. @pytest.mark.parametrize("mqtt_password", ["secret"])
  124. def test__run_authentication_missing_username(mqtt_host, mqtt_port, mqtt_password):
  125. with unittest.mock.patch("paho.mqtt.client.Client"):
  126. with pytest.raises(ValueError):
  127. switchbot_mqtt._run(
  128. mqtt_host=mqtt_host,
  129. mqtt_port=mqtt_port,
  130. mqtt_username=None,
  131. mqtt_password=mqtt_password,
  132. retry_count=3,
  133. device_passwords={},
  134. fetch_device_info=True,
  135. )
  136. def _mock_actor_class(
  137. command_topic_levels: typing.List[switchbot_mqtt._MQTTTopicLevel],
  138. ) -> typing.Type:
  139. class _ActorMock(switchbot_mqtt._MQTTControlledActor):
  140. MQTT_COMMAND_TOPIC_LEVELS = command_topic_levels
  141. def __init__(self, mac_address, retry_count, password):
  142. super().__init__(
  143. mac_address=mac_address, retry_count=retry_count, password=password
  144. )
  145. def execute_command(
  146. self,
  147. mqtt_message_payload: bytes,
  148. mqtt_client: Client,
  149. update_device_info: bool,
  150. ):
  151. pass
  152. return _ActorMock
  153. @pytest.mark.parametrize(
  154. ("command_topic_levels", "topic", "payload", "expected_mac_address"),
  155. [
  156. (
  157. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS,
  158. b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set",
  159. b"ON",
  160. "aa:bb:cc:dd:ee:ff",
  161. ),
  162. (
  163. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS,
  164. b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set",
  165. b"OFF",
  166. "aa:bb:cc:dd:ee:ff",
  167. ),
  168. (
  169. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS,
  170. b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set",
  171. b"on",
  172. "aa:bb:cc:dd:ee:ff",
  173. ),
  174. (
  175. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS,
  176. b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set",
  177. b"off",
  178. "aa:bb:cc:dd:ee:ff",
  179. ),
  180. (
  181. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS,
  182. b"homeassistant/switch/switchbot/aa:01:23:45:67:89/set",
  183. b"ON",
  184. "aa:01:23:45:67:89",
  185. ),
  186. (
  187. ["switchbot", switchbot_mqtt._MQTTTopicPlaceholder.MAC_ADDRESS],
  188. b"switchbot/aa:01:23:45:67:89",
  189. b"ON",
  190. "aa:01:23:45:67:89",
  191. ),
  192. (
  193. switchbot_mqtt._CurtainMotor.MQTT_COMMAND_TOPIC_LEVELS,
  194. b"homeassistant/cover/switchbot-curtain/aa:01:23:45:67:89/set",
  195. b"OPEN",
  196. "aa:01:23:45:67:89",
  197. ),
  198. ],
  199. )
  200. @pytest.mark.parametrize("retry_count", (3, 42))
  201. @pytest.mark.parametrize("fetch_device_info", [True, False])
  202. def test__mqtt_command_callback(
  203. caplog,
  204. command_topic_levels: typing.List[switchbot_mqtt._MQTTTopicLevel],
  205. topic: bytes,
  206. payload: bytes,
  207. expected_mac_address: str,
  208. retry_count: int,
  209. fetch_device_info: bool,
  210. ):
  211. ActorMock = _mock_actor_class(command_topic_levels)
  212. message = MQTTMessage(topic=topic)
  213. message.payload = payload
  214. callback_userdata = switchbot_mqtt._MQTTCallbackUserdata(
  215. retry_count=retry_count,
  216. device_passwords={},
  217. fetch_device_info=fetch_device_info,
  218. )
  219. with unittest.mock.patch.object(
  220. ActorMock, "__init__", return_value=None
  221. ) as init_mock, unittest.mock.patch.object(
  222. ActorMock, "execute_command"
  223. ) as execute_command_mock, caplog.at_level(
  224. logging.DEBUG
  225. ):
  226. ActorMock._mqtt_command_callback("client_dummy", callback_userdata, message)
  227. init_mock.assert_called_once_with(
  228. mac_address=expected_mac_address, retry_count=retry_count, password=None
  229. )
  230. execute_command_mock.assert_called_once_with(
  231. mqtt_client="client_dummy",
  232. mqtt_message_payload=payload,
  233. update_device_info=fetch_device_info,
  234. )
  235. assert caplog.record_tuples == [
  236. (
  237. "switchbot_mqtt",
  238. logging.DEBUG,
  239. "received topic={} payload={!r}".format(topic.decode(), payload),
  240. )
  241. ]
  242. @pytest.mark.parametrize(
  243. ("mac_address", "expected_password"),
  244. [
  245. ("11:22:33:44:55:66", None),
  246. ("aa:bb:cc:dd:ee:ff", "secret"),
  247. ("11:22:33:dd:ee:ff", "äöü"),
  248. ],
  249. )
  250. def test__mqtt_command_callback_password(mac_address, expected_password):
  251. ActorMock = _mock_actor_class(
  252. [
  253. "switchbot",
  254. switchbot_mqtt._MQTTTopicPlaceholder.MAC_ADDRESS,
  255. ]
  256. )
  257. message = MQTTMessage(topic=b"switchbot/" + mac_address.encode())
  258. message.payload = b"whatever"
  259. callback_userdata = switchbot_mqtt._MQTTCallbackUserdata(
  260. retry_count=3,
  261. device_passwords={
  262. "11:22:33:44:55:77": "test",
  263. "aa:bb:cc:dd:ee:ff": "secret",
  264. "11:22:33:dd:ee:ff": "äöü",
  265. },
  266. fetch_device_info=True,
  267. )
  268. with unittest.mock.patch.object(
  269. ActorMock, "__init__", return_value=None
  270. ) as init_mock, unittest.mock.patch.object(
  271. ActorMock, "execute_command"
  272. ) as execute_command_mock:
  273. ActorMock._mqtt_command_callback("client_dummy", callback_userdata, message)
  274. init_mock.assert_called_once_with(
  275. mac_address=mac_address, retry_count=3, password=expected_password
  276. )
  277. execute_command_mock.assert_called_once_with(
  278. mqtt_client="client_dummy",
  279. mqtt_message_payload=b"whatever",
  280. update_device_info=True,
  281. )
  282. @pytest.mark.parametrize(
  283. ("topic", "payload"),
  284. [
  285. (b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff", b"on"),
  286. (b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/change", b"ON"),
  287. (b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set/suffix", b"ON"),
  288. ],
  289. )
  290. def test__mqtt_command_callback_unexpected_topic(caplog, topic: bytes, payload: bytes):
  291. ActorMock = _mock_actor_class(
  292. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS
  293. )
  294. message = MQTTMessage(topic=topic)
  295. message.payload = payload
  296. with unittest.mock.patch.object(
  297. ActorMock, "__init__", return_value=None
  298. ) as init_mock, unittest.mock.patch.object(
  299. ActorMock, "execute_command"
  300. ) as execute_command_mock, caplog.at_level(
  301. logging.DEBUG
  302. ):
  303. ActorMock._mqtt_command_callback(
  304. "client_dummy",
  305. switchbot_mqtt._MQTTCallbackUserdata(
  306. retry_count=3, device_passwords={}, fetch_device_info=True
  307. ),
  308. message,
  309. )
  310. init_mock.assert_not_called()
  311. execute_command_mock.assert_not_called()
  312. assert caplog.record_tuples == [
  313. (
  314. "switchbot_mqtt",
  315. logging.DEBUG,
  316. "received topic={} payload={!r}".format(topic.decode(), payload),
  317. ),
  318. (
  319. "switchbot_mqtt",
  320. logging.WARNING,
  321. "unexpected topic {}".format(topic.decode()),
  322. ),
  323. ]
  324. @pytest.mark.parametrize(("mac_address", "payload"), [("aa:01:23:4E:RR:OR", b"ON")])
  325. def test__mqtt_command_callback_invalid_mac_address(
  326. caplog, mac_address: str, payload: bytes
  327. ):
  328. ActorMock = _mock_actor_class(
  329. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS
  330. )
  331. topic = "homeassistant/switch/switchbot/{}/set".format(mac_address).encode()
  332. message = MQTTMessage(topic=topic)
  333. message.payload = payload
  334. with unittest.mock.patch.object(
  335. ActorMock, "__init__", return_value=None
  336. ) as init_mock, unittest.mock.patch.object(
  337. ActorMock, "execute_command"
  338. ) as execute_command_mock, caplog.at_level(
  339. logging.DEBUG
  340. ):
  341. ActorMock._mqtt_command_callback(
  342. "client_dummy",
  343. switchbot_mqtt._MQTTCallbackUserdata(
  344. retry_count=3, device_passwords={}, fetch_device_info=True
  345. ),
  346. message,
  347. )
  348. init_mock.assert_not_called()
  349. execute_command_mock.assert_not_called()
  350. assert caplog.record_tuples == [
  351. (
  352. "switchbot_mqtt",
  353. logging.DEBUG,
  354. "received topic={} payload={!r}".format(topic.decode(), payload),
  355. ),
  356. (
  357. "switchbot_mqtt",
  358. logging.WARNING,
  359. "invalid mac address {}".format(mac_address),
  360. ),
  361. ]
  362. @pytest.mark.parametrize(
  363. ("topic", "payload"),
  364. [(b"homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/set", b"ON")],
  365. )
  366. def test__mqtt_command_callback_ignore_retained(caplog, topic: bytes, payload: bytes):
  367. ActorMock = _mock_actor_class(
  368. switchbot_mqtt._ButtonAutomator.MQTT_COMMAND_TOPIC_LEVELS
  369. )
  370. message = MQTTMessage(topic=topic)
  371. message.payload = payload
  372. message.retain = True
  373. with unittest.mock.patch.object(
  374. ActorMock, "__init__", return_value=None
  375. ) as init_mock, unittest.mock.patch.object(
  376. ActorMock, "execute_command"
  377. ) as execute_command_mock, caplog.at_level(
  378. logging.DEBUG
  379. ):
  380. ActorMock._mqtt_command_callback(
  381. "client_dummy",
  382. switchbot_mqtt._MQTTCallbackUserdata(
  383. retry_count=4, device_passwords={}, fetch_device_info=True
  384. ),
  385. message,
  386. )
  387. init_mock.assert_not_called()
  388. execute_command_mock.assert_not_called()
  389. assert caplog.record_tuples == [
  390. (
  391. "switchbot_mqtt",
  392. logging.DEBUG,
  393. "received topic={} payload={!r}".format(topic.decode(), payload),
  394. ),
  395. ("switchbot_mqtt", logging.INFO, "ignoring retained message"),
  396. ]
  397. @pytest.mark.parametrize(
  398. ("state_topic_levels", "mac_address", "expected_topic"),
  399. # https://www.home-assistant.io/docs/mqtt/discovery/#switches
  400. [
  401. (
  402. switchbot_mqtt._ButtonAutomator.MQTT_STATE_TOPIC_LEVELS,
  403. "aa:bb:cc:dd:ee:ff",
  404. "homeassistant/switch/switchbot/aa:bb:cc:dd:ee:ff/state",
  405. ),
  406. (
  407. ["switchbot", switchbot_mqtt._MQTTTopicPlaceholder.MAC_ADDRESS, "state"],
  408. "aa:bb:cc:dd:ee:gg",
  409. "switchbot/aa:bb:cc:dd:ee:gg/state",
  410. ),
  411. ],
  412. )
  413. @pytest.mark.parametrize("state", [b"ON", b"CLOSE"])
  414. @pytest.mark.parametrize("return_code", [MQTT_ERR_SUCCESS, MQTT_ERR_QUEUE_SIZE])
  415. def test__report_state(
  416. caplog,
  417. state_topic_levels: typing.List[switchbot_mqtt._MQTTTopicLevel],
  418. mac_address: str,
  419. expected_topic: str,
  420. state: bytes,
  421. return_code: int,
  422. ):
  423. # pylint: disable=too-many-arguments
  424. class _ActorMock(switchbot_mqtt._MQTTControlledActor):
  425. MQTT_STATE_TOPIC_LEVELS = state_topic_levels
  426. def __init__(self, mac_address, retry_count, password):
  427. super().__init__(
  428. mac_address=mac_address, retry_count=retry_count, password=password
  429. )
  430. def execute_command(
  431. self,
  432. mqtt_message_payload: bytes,
  433. mqtt_client: Client,
  434. update_device_info: bool,
  435. ):
  436. pass
  437. mqtt_client_mock = unittest.mock.MagicMock()
  438. mqtt_client_mock.publish.return_value.rc = return_code
  439. with caplog.at_level(logging.DEBUG):
  440. actor = _ActorMock(mac_address=mac_address, retry_count=3, password=None)
  441. actor.report_state(state=state, mqtt_client=mqtt_client_mock)
  442. mqtt_client_mock.publish.assert_called_once_with(
  443. topic=expected_topic, payload=state, retain=True
  444. )
  445. assert caplog.record_tuples[0] == (
  446. "switchbot_mqtt",
  447. logging.DEBUG,
  448. "publishing topic={} payload={!r}".format(expected_topic, state),
  449. )
  450. if return_code == MQTT_ERR_SUCCESS:
  451. assert not caplog.records[1:]
  452. else:
  453. assert caplog.record_tuples[1:] == [
  454. (
  455. "switchbot_mqtt",
  456. logging.ERROR,
  457. "Failed to publish MQTT message on topic {} (rc={})".format(
  458. expected_topic, return_code
  459. ),
  460. )
  461. ]