Convert MQTT to use asyncio (#115910)

This commit is contained in:
J. Nick Koston 2024-04-21 22:33:58 +02:00 committed by GitHub
parent 5a24690d79
commit 423544401e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 464 additions and 90 deletions

View file

@ -4,17 +4,22 @@ import asyncio
from copy import deepcopy
from datetime import datetime, timedelta
import json
import socket
import ssl
from typing import Any, TypedDict
from unittest.mock import ANY, MagicMock, call, mock_open, patch
from freezegun.api import FrozenDateTimeFactory
import paho.mqtt.client as paho_mqtt
import pytest
import voluptuous as vol
from homeassistant.components import mqtt
from homeassistant.components.mqtt import debug_info
from homeassistant.components.mqtt.client import EnsureJobAfterCooldown
from homeassistant.components.mqtt.client import (
RECONNECT_INTERVAL_SECONDS,
EnsureJobAfterCooldown,
)
from homeassistant.components.mqtt.mixins import MQTT_ENTITY_DEVICE_INFO_SCHEMA
from homeassistant.components.mqtt.models import (
MessageCallbackType,
@ -146,7 +151,7 @@ async def test_mqtt_disconnects_on_home_assistant_stop(
hass.bus.fire(EVENT_HOMEASSISTANT_STOP)
await hass.async_block_till_done()
await hass.async_block_till_done()
assert mqtt_client_mock.loop_stop.call_count == 1
assert mqtt_client_mock.disconnect.call_count == 1
async def test_mqtt_await_ack_at_disconnect(
@ -161,8 +166,14 @@ async def test_mqtt_await_ack_at_disconnect(
rc = 0
with patch("paho.mqtt.client.Client") as mock_client:
mock_client().connect = MagicMock(return_value=0)
mock_client().publish = MagicMock(return_value=FakeInfo())
mqtt_client = mock_client.return_value
mqtt_client.connect = MagicMock(
return_value=0,
side_effect=lambda *args, **kwargs: hass.loop.call_soon_threadsafe(
mqtt_client.on_connect, mqtt_client, None, 0, 0, 0
),
)
mqtt_client.publish = MagicMock(return_value=FakeInfo())
entry = MockConfigEntry(
domain=mqtt.DOMAIN,
data={"certificate": "auto", mqtt.CONF_BROKER: "test-broker"},
@ -1669,6 +1680,7 @@ async def test_not_calling_subscribe_when_unsubscribed_within_cooldown(
the subscribe cool down period has ended.
"""
mqtt_mock = await mqtt_mock_entry()
mqtt_client_mock.subscribe.reset_mock()
# Fake that the client is connected
mqtt_mock().connected = True
@ -1925,6 +1937,7 @@ async def test_canceling_debouncer_on_shutdown(
"""Test canceling the debouncer when HA shuts down."""
mqtt_mock = await mqtt_mock_entry()
mqtt_client_mock.subscribe.reset_mock()
# Fake that the client is connected
mqtt_mock().connected = True
@ -2008,7 +2021,7 @@ async def test_initial_setup_logs_error(
"""Test for setup failure if initial client connection fails."""
entry = MockConfigEntry(domain=mqtt.DOMAIN, data={mqtt.CONF_BROKER: "test-broker"})
entry.add_to_hass(hass)
mqtt_client_mock.connect.return_value = 1
mqtt_client_mock.connect.side_effect = MagicMock(return_value=1)
try:
assert await hass.config_entries.async_setup(entry.entry_id)
except HomeAssistantError:
@ -2230,7 +2243,12 @@ async def test_handle_mqtt_timeout_on_callback(
mock_client = mock_client.return_value
mock_client.publish.return_value = FakeInfo()
mock_client.subscribe.side_effect = _mock_ack
mock_client.connect.return_value = 0
mock_client.connect = MagicMock(
return_value=0,
side_effect=lambda *args, **kwargs: hass.loop.call_soon_threadsafe(
mock_client.on_connect, mock_client, None, 0, 0, 0
),
)
entry = MockConfigEntry(
domain=mqtt.DOMAIN, data={mqtt.CONF_BROKER: "test-broker"}
@ -4144,3 +4162,179 @@ async def test_multi_platform_discovery(
)
is not None
)
@patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
async def test_auto_reconnect(
hass: HomeAssistant,
mqtt_client_mock: MqttMockPahoClient,
mqtt_mock_entry: MqttMockHAClientGenerator,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test reconnection is automatically done."""
mqtt_mock = await mqtt_mock_entry()
await hass.async_block_till_done()
assert mqtt_mock.connected is True
mqtt_client_mock.reconnect.reset_mock()
mqtt_client_mock.disconnect()
mqtt_client_mock.on_disconnect(None, None, 0)
await hass.async_block_till_done()
mqtt_client_mock.reconnect.side_effect = OSError("foo")
async_fire_time_changed(
hass, utcnow() + timedelta(seconds=RECONNECT_INTERVAL_SECONDS)
)
await hass.async_block_till_done()
assert len(mqtt_client_mock.reconnect.mock_calls) == 1
assert "Error re-connecting to MQTT server due to exception: foo" in caplog.text
mqtt_client_mock.reconnect.side_effect = None
async_fire_time_changed(
hass, utcnow() + timedelta(seconds=RECONNECT_INTERVAL_SECONDS)
)
await hass.async_block_till_done()
assert len(mqtt_client_mock.reconnect.mock_calls) == 2
hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
mqtt_client_mock.disconnect()
mqtt_client_mock.on_disconnect(None, None, 0)
await hass.async_block_till_done()
async_fire_time_changed(
hass, utcnow() + timedelta(seconds=RECONNECT_INTERVAL_SECONDS)
)
await hass.async_block_till_done()
# Should not reconnect after stop
assert len(mqtt_client_mock.reconnect.mock_calls) == 2
@patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
async def test_server_sock_connect_and_disconnect(
hass: HomeAssistant,
mqtt_client_mock: MqttMockPahoClient,
mqtt_mock_entry: MqttMockHAClientGenerator,
calls: list[ReceiveMessage],
record_calls: MessageCallbackType,
) -> None:
"""Test handling the socket connected and disconnected."""
mqtt_mock = await mqtt_mock_entry()
await hass.async_block_till_done()
assert mqtt_mock.connected is True
mqtt_client_mock.loop_misc.return_value = paho_mqtt.MQTT_ERR_SUCCESS
client, server = socket.socketpair(
family=socket.AF_UNIX, type=socket.SOCK_STREAM, proto=0
)
client.setblocking(False)
server.setblocking(False)
mqtt_client_mock.on_socket_open(mqtt_client_mock, None, client)
mqtt_client_mock.on_socket_register_write(mqtt_client_mock, None, client)
await hass.async_block_till_done()
server.close() # mock the server closing the connection on us
unsub = await mqtt.async_subscribe(hass, "test-topic", record_calls)
mqtt_client_mock.loop_misc.return_value = paho_mqtt.MQTT_ERR_CONN_LOST
mqtt_client_mock.on_socket_unregister_write(mqtt_client_mock, None, client)
mqtt_client_mock.on_socket_close(mqtt_client_mock, None, client)
mqtt_client_mock.on_disconnect(mqtt_client_mock, None, client)
await hass.async_block_till_done()
unsub()
# Should have failed
assert len(calls) == 0
@patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
async def test_client_sock_failure_after_connect(
hass: HomeAssistant,
mqtt_client_mock: MqttMockPahoClient,
mqtt_mock_entry: MqttMockHAClientGenerator,
calls: list[ReceiveMessage],
record_calls: MessageCallbackType,
) -> None:
"""Test handling the socket connected and disconnected."""
mqtt_mock = await mqtt_mock_entry()
# Fake that the client is connected
mqtt_mock().connected = True
await hass.async_block_till_done()
assert mqtt_mock.connected is True
mqtt_client_mock.loop_misc.return_value = paho_mqtt.MQTT_ERR_SUCCESS
client, server = socket.socketpair(
family=socket.AF_UNIX, type=socket.SOCK_STREAM, proto=0
)
client.setblocking(False)
server.setblocking(False)
mqtt_client_mock.on_socket_open(mqtt_client_mock, None, client)
mqtt_client_mock.on_socket_register_writer(mqtt_client_mock, None, client)
await hass.async_block_till_done()
mqtt_client_mock.loop_write.side_effect = OSError("foo")
client.close() # close the client socket out from under the client
assert mqtt_mock.connected is True
unsub = await mqtt.async_subscribe(hass, "test-topic", record_calls)
async_fire_time_changed(hass, utcnow() + timedelta(seconds=5))
await hass.async_block_till_done()
unsub()
# Should have failed
assert len(calls) == 0
@patch("homeassistant.components.mqtt.client.INITIAL_SUBSCRIBE_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.DISCOVERY_COOLDOWN", 0.0)
@patch("homeassistant.components.mqtt.client.SUBSCRIBE_COOLDOWN", 0.0)
async def test_loop_write_failure(
hass: HomeAssistant,
mqtt_client_mock: MqttMockPahoClient,
mqtt_mock_entry: MqttMockHAClientGenerator,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test handling the socket connected and disconnected."""
mqtt_mock = await mqtt_mock_entry()
await hass.async_block_till_done()
assert mqtt_mock.connected is True
mqtt_client_mock.loop_misc.return_value = paho_mqtt.MQTT_ERR_SUCCESS
client, server = socket.socketpair(
family=socket.AF_UNIX, type=socket.SOCK_STREAM, proto=0
)
client.setblocking(False)
server.setblocking(False)
mqtt_client_mock.on_socket_open(mqtt_client_mock, None, client)
mqtt_client_mock.on_socket_register_write(mqtt_client_mock, None, client)
mqtt_client_mock.loop_write.return_value = paho_mqtt.MQTT_ERR_CONN_LOST
mqtt_client_mock.loop_read.return_value = paho_mqtt.MQTT_ERR_CONN_LOST
# Fill up the outgoing buffer to ensure that loop_write
# and loop_read are called that next time control is
# returned to the event loop
try:
for _ in range(1000):
server.send(b"long" * 100)
except BlockingIOError:
pass
server.close()
# Once for the reader callback
await hass.async_block_till_done()
# Another for the writer callback
await hass.async_block_till_done()
# Final for the disconnect callback
await hass.async_block_till_done()
assert "Disconnected from MQTT server mock-broker:1883 (7)" in caplog.text