jeans/core
jeans/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2025-04-24 01:08:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Upgrade paho-mqtt API to v2 (#137613)

Browse Source 
* Upgrade paho-mqtt API to v2

* Refactor on_connect callback

* Add tests

* Fix Tasmota tests
This commit is contained in:
Jan Bouwhuis 2025-02-13 22:13:19 +01:00 committed by GitHub
parent bbbad90ca2
commit d6b7762dd6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 171 additions and 93 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
homeassistant/components/mqtt/async_client.py
View File

@ -6,7 +6,14 @@ from functools import lru_cache
from types import TracebackType
from typing import Self
from paho.mqtt.client import Client as MQTTClient
from paho.mqtt.client import (
CallbackOnConnect_v2,
CallbackOnDisconnect_v2,
CallbackOnPublish_v2,
CallbackOnSubscribe_v2,
CallbackOnUnsubscribe_v2,
Client as MQTTClient,
)
_MQTT_LOCK_COUNT = 7
@ -44,6 +51,12 @@ class AsyncMQTTClient(MQTTClient):
that is not needed since we are running in an async event loop.
"""
on_connect: CallbackOnConnect_v2
on_disconnect: CallbackOnDisconnect_v2
on_publish: CallbackOnPublish_v2
on_subscribe: CallbackOnSubscribe_v2
on_unsubscribe: CallbackOnUnsubscribe_v2
def setup(self) -> None:
"""Set up the client.

86
homeassistant/components/mqtt/client.py
View File

@ -311,8 +311,8 @@ class MqttClientSetup:
client_id = None
transport: str = config.get(CONF_TRANSPORT, DEFAULT_TRANSPORT)
self._client = AsyncMQTTClient(
mqtt.CallbackAPIVersion.VERSION1,
client_id,
callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
client_id=client_id,
protocol=proto,
transport=transport, # type: ignore[arg-type]
reconnect_on_failure=False,
@ -476,9 +476,9 @@ class MQTT:
mqttc.on_connect = self._async_mqtt_on_connect
mqttc.on_disconnect = self._async_mqtt_on_disconnect
mqttc.on_message = self._async_mqtt_on_message
mqttc.on_publish = self._async_mqtt_on_callback
mqttc.on_subscribe = self._async_mqtt_on_callback
mqttc.on_unsubscribe = self._async_mqtt_on_callback
mqttc.on_publish = self._async_mqtt_on_publish
mqttc.on_subscribe = self._async_mqtt_on_subscribe_unsubscribe
mqttc.on_unsubscribe = self._async_mqtt_on_subscribe_unsubscribe
# suppress exceptions at callback
mqttc.suppress_exceptions = True
@ -498,7 +498,7 @@ class MQTT:
def _async_reader_callback(self, client: mqtt.Client) -> None:
"""Handle reading data from the socket."""
if (status := client.loop_read(MAX_PACKETS_TO_READ)) != 0:
self._async_on_disconnect(status)
self._async_handle_callback_exception(status)
@callback
def _async_start_misc_periodic(self) -> None:
@ -593,7 +593,7 @@ class MQTT:
def _async_writer_callback(self, client: mqtt.Client) -> None:
"""Handle writing data to the socket."""
if (status := client.loop_write()) != 0:
self._async_on_disconnect(status)
self._async_handle_callback_exception(status)
def _on_socket_register_write(
self, client: mqtt.Client, userdata: Any, sock: SocketType
@ -983,9 +983,9 @@ class MQTT:
self,
_mqttc: mqtt.Client,
_userdata: None,
_flags: dict[str, int],
result_code: int,
properties: mqtt.Properties | None = None,
_connect_flags: mqtt.ConnectFlags,
reason_code: mqtt.ReasonCode,
_properties: mqtt.Properties | None = None,
) -> None:
"""On connect callback.
@ -993,19 +993,20 @@ class MQTT:
message.
"""
# pylint: disable-next=import-outside-toplevel
import paho.mqtt.client as mqtt
if result_code != mqtt.CONNACK_ACCEPTED:
if result_code in (
mqtt.CONNACK_REFUSED_BAD_USERNAME_PASSWORD,
mqtt.CONNACK_REFUSED_NOT_AUTHORIZED,
):
if reason_code.is_failure:
# 24: Continue authentication
# 25: Re-authenticate
# 134: Bad user name or password
# 135: Not authorized
# 140: Bad authentication method
if reason_code.value in (24, 25, 134, 135, 140):
self._should_reconnect = False
self.hass.async_create_task(self.async_disconnect())
self.config_entry.async_start_reauth(self.hass)
_LOGGER.error(
"Unable to connect to the MQTT broker: %s",
mqtt.connack_string(result_code),
reason_code.getName(), # type: ignore[no-untyped-call]
)
self._async_connection_result(False)
return
@ -1016,7 +1017,7 @@ class MQTT:
"Connected to MQTT server %s:%s (%s)",
self.conf[CONF_BROKER],
self.conf.get(CONF_PORT, DEFAULT_PORT),
result_code,
reason_code,
)
birth: dict[str, Any]
@ -1153,18 +1154,32 @@ class MQTT:
self._mqtt_data.state_write_requests.process_write_state_requests(msg)
@callback
def _async_mqtt_on_callback(
def _async_mqtt_on_publish(
self,
_mqttc: mqtt.Client,
_userdata: None,
mid: int,
_granted_qos_reason: tuple[int, ...] | mqtt.ReasonCodes | None = None,
_properties_reason: mqtt.ReasonCodes | None = None,
_reason_code: mqtt.ReasonCode,
_properties: mqtt.Properties | None,
) -> None:
"""Publish callback."""
self._async_mqtt_on_callback(mid)
@callback
def _async_mqtt_on_subscribe_unsubscribe(
self,
_mqttc: mqtt.Client,
_userdata: None,
mid: int,
_reason_code: list[mqtt.ReasonCode],
_properties: mqtt.Properties | None,
) -> None:
"""Subscribe / Unsubscribe callback."""
self._async_mqtt_on_callback(mid)
@callback
def _async_mqtt_on_callback(self, mid: int) -> None:
"""Publish / Subscribe / Unsubscribe callback."""
# The callback signature for on_unsubscribe is different from on_subscribe
# see https://github.com/eclipse/paho.mqtt.python/issues/687
# properties and reason codes are not used in Home Assistant
future = self._async_get_mid_future(mid)
if future.done() and (future.cancelled() or future.exception()):
# Timed out or cancelled
@ -1180,19 +1195,28 @@ class MQTT:
self._pending_operations[mid] = future
return future
@callback
def _async_handle_callback_exception(self, status: mqtt.MQTTErrorCode) -> None:
"""Handle a callback exception."""
# We don't import on the top because some integrations
# should be able to optionally rely on MQTT.
import paho.mqtt.client as mqtt # pylint: disable=import-outside-toplevel
_LOGGER.warning(
"Error returned from MQTT server: %s",
mqtt.error_string(status),
)
@callback
def _async_mqtt_on_disconnect(
self,
_mqttc: mqtt.Client,
_userdata: None,
result_code: int,
_disconnect_flags: mqtt.DisconnectFlags,
reason_code: mqtt.ReasonCode,
properties: mqtt.Properties | None = None,
) -> None:
"""Disconnected callback."""
self._async_on_disconnect(result_code)
@callback
def _async_on_disconnect(self, result_code: int) -> None:
if not self.connected:
# This function is re-entrant and may be called multiple times
# when there is a broken pipe error.
@ -1203,11 +1227,11 @@ class MQTT:
self.connected = False
async_dispatcher_send(self.hass, MQTT_CONNECTION_STATE, False)
_LOGGER.log(
logging.INFO if result_code == 0 else logging.DEBUG,
logging.INFO if reason_code == 0 else logging.DEBUG,
"Disconnected from MQTT server %s:%s (%s)",
self.conf[CONF_BROKER],
self.conf.get(CONF_PORT, DEFAULT_PORT),
result_code,
reason_code,
)
@callback

12
homeassistant/components/mqtt/config_flow.py
View File

@ -1023,14 +1023,14 @@ def try_connection(
result: queue.Queue[bool] = queue.Queue(maxsize=1)
def on_connect(
client_: mqtt.Client,
userdata: None,
flags: dict[str, Any],
result_code: int,
properties: mqtt.Properties | None = None,
_mqttc: mqtt.Client,
_userdata: None,
_connect_flags: mqtt.ConnectFlags,
reason_code: mqtt.ReasonCode,
_properties: mqtt.Properties | None = None,
) -> None:
"""Handle connection result."""
result.put(result_code == mqtt.CONNACK_ACCEPTED)
result.put(not reason_code.is_failure)
client.on_connect = on_connect

19
tests/common.py
View File

@ -410,6 +410,25 @@ def async_mock_intent(hass: HomeAssistant, intent_typ: str) -> list[intent.Inten
return intents
class MockMqttReasonCode:
"""Class to fake a MQTT ReasonCode."""
value: int
is_failure: bool
def __init__(
self, value: int = 0, is_failure: bool = False, name: str = "Success"
) -> None:
"""Initialize the mock reason code."""
self.value = value
self.is_failure = is_failure
self._name = name
def getName(self) -> str:
"""Return the name of the reason code."""
return self._name
@callback
def async_fire_mqtt_message(
hass: HomeAssistant,

83
tests/components/mqtt/test_client.py
View File

@ -32,6 +32,7 @@ from .test_common import help_all_subscribe_calls
from tests.common import (
MockConfigEntry,
MockMqttReasonCode,
async_fire_mqtt_message,
async_fire_time_changed,
)
@ -94,7 +95,7 @@ async def test_mqtt_await_ack_at_disconnect(hass: HomeAssistant) -> None:
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.on_connect, mqtt_client, None, 0, MockMqttReasonCode()
),
)
mqtt_client.publish = MagicMock(return_value=FakeInfo())
@ -119,7 +120,7 @@ async def test_mqtt_await_ack_at_disconnect(hass: HomeAssistant) -> None:
)
await asyncio.sleep(0)
# Simulate late ACK callback from client with mid 100
mqtt_client.on_publish(0, 0, 100)
mqtt_client.on_publish(0, 0, 100, MockMqttReasonCode(), None)
# disconnect the MQTT client
await hass.async_stop()
await hass.async_block_till_done()
@ -778,10 +779,10 @@ async def test_replaying_payload_same_topic(
calls_a = []
calls_b = []
mqtt_client_mock.reset_mock()
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await mock_debouncer.wait()
mqtt_client_mock.subscribe.assert_called()
# Simulate a (retained) message played back after reconnecting
@ -908,10 +909,10 @@ async def test_replaying_payload_wildcard_topic(
calls_a = []
calls_b = []
mqtt_client_mock.reset_mock()
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await mock_debouncer.wait()
mqtt_client_mock.subscribe.assert_called()
@ -1045,7 +1046,7 @@ async def test_restore_subscriptions_on_reconnect(
assert ("test/state", 0) in help_all_subscribe_calls(mqtt_client_mock)
mqtt_client_mock.reset_mock()
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
# Test to subscribe orther topic while the client is not connected
await mqtt.async_subscribe(hass, "test/other", record_calls)
@ -1053,7 +1054,7 @@ async def test_restore_subscriptions_on_reconnect(
assert ("test/other", 0) not in help_all_subscribe_calls(mqtt_client_mock)
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await mock_debouncer.wait()
# Assert all subscriptions are performed at the broker
assert ("test/state", 0) in help_all_subscribe_calls(mqtt_client_mock)
@ -1089,10 +1090,10 @@ async def test_restore_all_active_subscriptions_on_reconnect(
unsub()
assert mqtt_client_mock.unsubscribe.call_count == 0
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
# wait for cooldown
await mock_debouncer.wait()
@ -1160,27 +1161,37 @@ async def test_logs_error_if_no_connect_broker(
) -> None:
"""Test for setup failure if connection to broker is missing."""
mqtt_client_mock = setup_with_birth_msg_client_mock
# test with rc = 3 -> broker unavailable
mqtt_client_mock.on_disconnect(Mock(), None, 0)
mqtt_client_mock.on_connect(Mock(), None, None, 3)
await hass.async_block_till_done()
assert (
"Unable to connect to the MQTT broker: Connection Refused: broker unavailable."
in caplog.text
# test with reason code = 136 -> server unavailable
mqtt_client_mock.on_disconnect(Mock(), None, None, MockMqttReasonCode())
mqtt_client_mock.on_connect(
Mock(),
None,
None,
MockMqttReasonCode(value=136, is_failure=True, name="Server unavailable"),
)
await hass.async_block_till_done()
assert "Unable to connect to the MQTT broker: Server unavailable" in caplog.text
@pytest.mark.parametrize("return_code", [4, 5])
@pytest.mark.parametrize(
"reason_code",
[
MockMqttReasonCode(
value=134, is_failure=True, name="Bad user name or password"
),
MockMqttReasonCode(value=135, is_failure=True, name="Not authorized"),
],
)
async def test_triggers_reauth_flow_if_auth_fails(
hass: HomeAssistant,
setup_with_birth_msg_client_mock: MqttMockPahoClient,
return_code: int,
reason_code: MockMqttReasonCode,
) -> None:
"""Test re-auth is triggered if authentication is failing."""
mqtt_client_mock = setup_with_birth_msg_client_mock
# test with rc = 4 -> CONNACK_REFUSED_NOT_AUTHORIZED and 5 -> CONNACK_REFUSED_BAD_USERNAME_PASSWORD
mqtt_client_mock.on_disconnect(Mock(), None, 0)
mqtt_client_mock.on_connect(Mock(), None, None, return_code)
mqtt_client_mock.on_disconnect(Mock(), None, 0, MockMqttReasonCode(), None)
mqtt_client_mock.on_connect(Mock(), None, None, reason_code)
await hass.async_block_till_done()
flows = hass.config_entries.flow.async_progress()
assert len(flows) == 1
@ -1197,7 +1208,9 @@ async def test_handle_mqtt_on_callback(
mqtt_client_mock = setup_with_birth_msg_client_mock
with patch.object(mqtt_client_mock, "get_mid", return_value=100):
# Simulate an ACK for mid == 100, this will call mqtt_mock._async_get_mid_future(mid)
mqtt_client_mock.on_publish(mqtt_client_mock, None, 100)
mqtt_client_mock.on_publish(
mqtt_client_mock, None, 100, MockMqttReasonCode(), None
)
await hass.async_block_till_done()
# Make sure the ACK has been received
await hass.async_block_till_done()
@ -1219,7 +1232,7 @@ async def test_handle_mqtt_on_callback_after_cancellation(
# Simulate the mid future getting a cancellation
mqtt_mock()._async_get_mid_future(101).cancel()
# Simulate an ACK for mid == 101, being received after the cancellation
mqtt_client_mock.on_publish(mqtt_client_mock, None, 101)
mqtt_client_mock.on_publish(mqtt_client_mock, None, 101, MockMqttReasonCode(), None)
await hass.async_block_till_done()
assert "No ACK from MQTT server" not in caplog.text
assert "InvalidStateError" not in caplog.text
@ -1236,7 +1249,7 @@ async def test_handle_mqtt_on_callback_after_timeout(
# Simulate the mid future getting a timeout
mqtt_mock()._async_get_mid_future(101).set_exception(asyncio.TimeoutError)
# Simulate an ACK for mid == 101, being received after the timeout
mqtt_client_mock.on_publish(mqtt_client_mock, None, 101)
mqtt_client_mock.on_publish(mqtt_client_mock, None, 101, MockMqttReasonCode(), None)
await hass.async_block_till_done()
assert "No ACK from MQTT server" not in caplog.text
assert "InvalidStateError" not in caplog.text
@ -1388,7 +1401,7 @@ async def test_handle_mqtt_timeout_on_callback(
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
mock_client.on_connect, mock_client, None, 0, MockMqttReasonCode()
),
)
@ -1777,12 +1790,12 @@ async def test_mqtt_subscribes_topics_on_connect(
await mqtt.async_subscribe(hass, "still/pending", record_calls, 1)
await mock_debouncer.wait()
mqtt_client_mock.on_disconnect(Mock(), None, 0)
mqtt_client_mock.on_disconnect(Mock(), None, 0, MockMqttReasonCode())
mqtt_client_mock.reset_mock()
mock_debouncer.clear()
mqtt_client_mock.on_connect(Mock(), None, 0, 0)
mqtt_client_mock.on_connect(Mock(), None, 0, MockMqttReasonCode())
await mock_debouncer.wait()
subscribe_calls = help_all_subscribe_calls(mqtt_client_mock)
@ -1837,12 +1850,12 @@ async def test_mqtt_subscribes_wildcard_topics_in_correct_order(
# Assert the initial wildcard topic subscription order
_assert_subscription_order()
mqtt_client_mock.on_disconnect(Mock(), None, 0)
mqtt_client_mock.on_disconnect(Mock(), None, 0, MockMqttReasonCode())
mqtt_client_mock.reset_mock()
mock_debouncer.clear()
mqtt_client_mock.on_connect(Mock(), None, 0, 0)
mqtt_client_mock.on_connect(Mock(), None, 0, MockMqttReasonCode())
await mock_debouncer.wait()
# Assert the wildcard topic subscription order after a reconnect
@ -1868,12 +1881,12 @@ async def test_mqtt_discovery_not_subscribes_when_disabled(
assert (f"homeassistant/{component}/+/config", 0) not in subscribe_calls
assert (f"homeassistant/{component}/+/+/config", 0) not in subscribe_calls
mqtt_client_mock.on_disconnect(Mock(), None, 0)
mqtt_client_mock.on_disconnect(Mock(), None, 0, MockMqttReasonCode())
mqtt_client_mock.reset_mock()
mock_debouncer.clear()
mqtt_client_mock.on_connect(Mock(), None, 0, 0)
mqtt_client_mock.on_connect(Mock(), None, 0, MockMqttReasonCode())
await mock_debouncer.wait()
subscribe_calls = help_all_subscribe_calls(mqtt_client_mock)
@ -1968,7 +1981,7 @@ async def test_auto_reconnect(
mqtt_client_mock.reconnect.reset_mock()
mqtt_client_mock.disconnect()
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
mqtt_client_mock.reconnect.side_effect = exception("foo")
@ -1989,7 +2002,7 @@ async def test_auto_reconnect(
hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
mqtt_client_mock.disconnect()
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
async_fire_time_changed(
@ -2031,7 +2044,7 @@ async def test_server_sock_connect_and_disconnect(
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)
mqtt_client_mock.on_disconnect(mqtt_client_mock, None, None, MockMqttReasonCode())
await hass.async_block_till_done()
mock_debouncer.clear()
unsub()
@ -2169,4 +2182,4 @@ async def test_loop_write_failure(
# Final for the disconnect callback
await hass.async_block_till_done()
assert "Disconnected from MQTT server test-broker:1883" in caplog.text
assert "Error returned from MQTT server: The connection was lost." in caplog.text

8
tests/components/mqtt/test_config_flow.py
View File

@ -28,7 +28,7 @@ from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType
from homeassistant.helpers.service_info.hassio import HassioServiceInfo
from tests.common import MockConfigEntry
from tests.common import MockConfigEntry, MockMqttReasonCode
from tests.typing import MqttMockHAClientGenerator, MqttMockPahoClient
ADD_ON_DISCOVERY_INFO = {
@ -143,16 +143,16 @@ def mock_try_connection_success() -> Generator[MqttMockPahoClient]:
def loop_start():
"""Simulate connect on loop start."""
mock_client().on_connect(mock_client, None, None, 0)
mock_client().on_connect(mock_client, None, None, MockMqttReasonCode(), None)
def _subscribe(topic, qos=0):
mid = get_mid()
mock_client().on_subscribe(mock_client, 0, mid)
mock_client().on_subscribe(mock_client, 0, mid, [MockMqttReasonCode()], None)
return (0, mid)
def _unsubscribe(topic):
mid = get_mid()
mock_client().on_unsubscribe(mock_client, 0, mid)
mock_client().on_unsubscribe(mock_client, 0, mid, [MockMqttReasonCode()], None)
return (0, mid)
with patch(

10
tests/components/mqtt/test_init.py
View File

@ -45,6 +45,7 @@ from tests.common import (
MockConfigEntry,
MockEntity,
MockEntityPlatform,
MockMqttReasonCode,
async_fire_mqtt_message,
async_fire_time_changed,
mock_restore_cache,
@ -1572,6 +1573,7 @@ async def test_subscribe_connection_status(
setup_with_birth_msg_client_mock: MqttMockPahoClient,
) -> None:
"""Test connextion status subscription."""
mqtt_client_mock = setup_with_birth_msg_client_mock
mqtt_connected_calls_callback: list[bool] = []
mqtt_connected_calls_async: list[bool] = []
@ -1589,7 +1591,7 @@ async def test_subscribe_connection_status(
assert mqtt.is_connected(hass) is True
# Mock disconnect status
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
assert mqtt.is_connected(hass) is False
@ -1603,12 +1605,12 @@ async def test_subscribe_connection_status(
# Mock connect status
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, 0, 0)
mqtt_client_mock.on_connect(None, None, 0, MockMqttReasonCode())
await mock_debouncer.wait()
assert mqtt.is_connected(hass) is True
# Mock disconnect status
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
assert mqtt.is_connected(hass) is False
@ -1618,7 +1620,7 @@ async def test_subscribe_connection_status(
# Mock connect status
mock_debouncer.clear()
mqtt_client_mock.on_connect(None, None, 0, 0)
mqtt_client_mock.on_connect(None, None, 0, MockMqttReasonCode())
await mock_debouncer.wait()
assert mqtt.is_connected(hass) is True

14
tests/components/tasmota/test_common.py
View File

@ -27,7 +27,7 @@ from homeassistant.const import STATE_UNAVAILABLE
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from tests.common import async_fire_mqtt_message
from tests.common import MockMqttReasonCode, async_fire_mqtt_message
from tests.typing import MqttMockHAClient, MqttMockPahoClient, WebSocketGenerator
DEFAULT_CONFIG = {
@ -165,7 +165,7 @@ async def help_test_availability_when_connection_lost(
# Disconnected from MQTT server -> state changed to unavailable
mqtt_mock.connected = False
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -174,7 +174,7 @@ async def help_test_availability_when_connection_lost(
# Reconnected to MQTT server -> state still unavailable
mqtt_mock.connected = True
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -226,7 +226,7 @@ async def help_test_deep_sleep_availability_when_connection_lost(
# Disconnected from MQTT server -> state changed to unavailable
mqtt_mock.connected = False
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -235,7 +235,7 @@ async def help_test_deep_sleep_availability_when_connection_lost(
# Reconnected to MQTT server -> state no longer unavailable
mqtt_mock.connected = True
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -478,7 +478,7 @@ async def help_test_availability_poll_state(
# Disconnected from MQTT server
mqtt_mock.connected = False
mqtt_client_mock.on_disconnect(None, None, 0)
mqtt_client_mock.on_disconnect(None, None, 0, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()
@ -486,7 +486,7 @@ async def help_test_availability_poll_state(
# Reconnected to MQTT server
mqtt_mock.connected = True
mqtt_client_mock.on_connect(None, None, None, 0)
mqtt_client_mock.on_connect(None, None, None, MockMqttReasonCode())
await hass.async_block_till_done()
await hass.async_block_till_done()
await hass.async_block_till_done()

17
tests/conftest.py
View File

@ -118,6 +118,7 @@ from .common import ( # noqa: E402, isort:skip
CLIENT_ID,
INSTANCES,
MockConfigEntry,
MockMqttReasonCode,
MockUser,
async_fire_mqtt_message,
async_test_home_assistant,
@ -969,17 +970,23 @@ def mqtt_client_mock(hass: HomeAssistant) -> Generator[MqttMockPahoClient]:
def _async_fire_mqtt_message(topic, payload, qos, retain):
async_fire_mqtt_message(hass, topic, payload or b"", qos, retain)
mid = get_mid()
hass.loop.call_soon(mock_client.on_publish, 0, 0, mid)
hass.loop.call_soon(
mock_client.on_publish, Mock(), 0, mid, MockMqttReasonCode(), None
)
return FakeInfo(mid)
def _subscribe(topic, qos=0):
mid = get_mid()
hass.loop.call_soon(mock_client.on_subscribe, 0, 0, mid)
hass.loop.call_soon(
mock_client.on_subscribe, Mock(), 0, mid, [MockMqttReasonCode()], None
)
return (0, mid)
def _unsubscribe(topic):
mid = get_mid()
hass.loop.call_soon(mock_client.on_unsubscribe, 0, 0, mid)
hass.loop.call_soon(
mock_client.on_unsubscribe, Mock(), 0, mid, [MockMqttReasonCode()], None
)
return (0, mid)
def _connect(*args, **kwargs):
@ -988,7 +995,7 @@ def mqtt_client_mock(hass: HomeAssistant) -> Generator[MqttMockPahoClient]:
# the behavior.
mock_client.reconnect()
hass.loop.call_soon_threadsafe(
mock_client.on_connect, mock_client, None, 0, 0, 0
mock_client.on_connect, mock_client, None, 0, MockMqttReasonCode()
)
mock_client.on_socket_open(
mock_client, None, Mock(fileno=Mock(return_value=-1))
@ -1065,7 +1072,7 @@ async def _mqtt_mock_entry(
# connected set to True to get a more realistic behavior when subscribing
mock_mqtt_instance.connected = True
mqtt_client_mock.on_connect(mqtt_client_mock, None, 0, 0, 0)
mqtt_client_mock.on_connect(mqtt_client_mock, None, 0, MockMqttReasonCode())
async_dispatcher_send(hass, mqtt.MQTT_CONNECTION_STATE, True)
await hass.async_block_till_done()