aioesphomeapi/tests/test_reconnect_logic.py
2023-12-12 07:24:31 -10:00

814 lines
24 KiB
Python

from __future__ import annotations
import asyncio
import logging
from functools import partial
from ipaddress import ip_address
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from zeroconf import (
DNSAddress,
DNSPointer,
DNSRecord,
RecordUpdate,
Zeroconf,
current_time_millis,
)
from zeroconf.asyncio import AsyncZeroconf
from zeroconf.const import _CLASS_IN, _TYPE_A, _TYPE_PTR
from aioesphomeapi import APIConnectionError, RequiresEncryptionAPIError
from aioesphomeapi._frame_helper.plain_text import APIPlaintextFrameHelper
from aioesphomeapi.client import APIClient
from aioesphomeapi.reconnect_logic import (
MAXIMUM_BACKOFF_TRIES,
ReconnectLogic,
ReconnectLogicState,
)
from .common import (
get_mock_async_zeroconf,
get_mock_zeroconf,
mock_data_received,
send_plaintext_connect_response,
send_plaintext_hello,
)
from .conftest import _create_mock_transport_protocol
logging.getLogger("aioesphomeapi").setLevel(logging.DEBUG)
async def slow_connect_fail(*args, **kwargs):
await asyncio.sleep(10)
raise APIConnectionError
async def quick_connect_fail(*args, **kwargs):
raise APIConnectionError
@pytest.mark.asyncio
async def test_reconnect_logic_name_from_host():
"""Test that the name is set correctly from the host."""
cli = APIClient(
address="mydevice.local",
port=6052,
password=None,
)
async def on_disconnect(expected_disconnect: bool) -> None:
pass
async def on_connect() -> None:
pass
ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=MagicMock(spec=AsyncZeroconf),
)
assert cli.log_name == "mydevice.local"
@pytest.mark.asyncio
async def test_reconnect_logic_name_from_host_and_set():
"""Test that the name is set correctly from the host."""
cli = APIClient(
address="mydevice.local",
port=6052,
password=None,
)
async def on_disconnect(expected_disconnect: bool) -> None:
pass
async def on_connect() -> None:
pass
rl = ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
)
assert rl.name == "mydevice"
assert cli.log_name == "mydevice.local"
@pytest.mark.asyncio
async def test_reconnect_logic_name_from_address():
"""Test that the name is set correctly from the address."""
cli = APIClient(
address="1.2.3.4",
port=6052,
password=None,
)
async def on_disconnect(expected_disconnect: bool) -> None:
pass
async def on_connect() -> None:
pass
ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
)
assert cli.log_name == "1.2.3.4"
@pytest.mark.asyncio
async def test_reconnect_logic_name_from_name():
"""Test that the name is set correctly from the address."""
cli = APIClient(
address="1.2.3.4",
port=6052,
password=None,
)
async def on_disconnect(expected_disconnect: bool) -> None:
pass
async def on_connect() -> None:
pass
ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
)
assert cli.log_name == "mydevice @ 1.2.3.4"
@pytest.mark.asyncio
async def test_reconnect_logic_name_from_cli_address():
"""Test that the name is set correctly from the address."""
cli = APIClient(
address="mydevice",
port=6052,
password=None,
)
async def on_disconnect(expected_disconnect: bool) -> None:
pass
async def on_connect() -> None:
pass
rl = ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
)
assert cli.log_name == "mydevice"
assert rl.name == "mydevice"
@pytest.mark.asyncio
async def test_reconnect_logic_state(patchable_api_client: APIClient):
"""Test that reconnect logic state changes."""
on_disconnect_called = []
on_connect_called = []
on_connect_fail_called = []
cli = patchable_api_client
async def on_disconnect(expected_disconnect: bool) -> None:
nonlocal on_disconnect_called
on_disconnect_called.append(expected_disconnect)
async def on_connect() -> None:
nonlocal on_connect_called
on_connect_called.append(True)
async def on_connect_fail(connect_exception: Exception) -> None:
nonlocal on_connect_called
on_connect_fail_called.append(connect_exception)
rl = ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
on_connect_error=on_connect_fail,
)
assert cli.log_name == "mydevice @ 1.2.3.4"
with patch.object(cli, "start_connection", side_effect=APIConnectionError):
await rl.start()
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 0
assert len(on_connect_fail_called) == 1
assert isinstance(on_connect_fail_called[-1], APIConnectionError)
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert rl._tries == 1
with patch.object(cli, "start_connection"), patch.object(
cli, "finish_connection", side_effect=RequiresEncryptionAPIError
):
await rl.start()
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 0
assert len(on_connect_fail_called) == 2
assert isinstance(on_connect_fail_called[-1], RequiresEncryptionAPIError)
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert rl._tries == MAXIMUM_BACKOFF_TRIES
with patch.object(cli, "start_connection"), patch.object(cli, "finish_connection"):
await rl.start()
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 1
assert len(on_connect_fail_called) == 2
assert rl._connection_state is ReconnectLogicState.READY
assert rl._tries == 0
await rl.stop()
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_reconnect_retry(
patchable_api_client: APIClient, caplog: pytest.LogCaptureFixture
):
"""Test that reconnect logic retry."""
on_disconnect_called = []
on_connect_called = []
on_connect_fail_called = []
cli = patchable_api_client
async def on_disconnect(expected_disconnect: bool) -> None:
nonlocal on_disconnect_called
on_disconnect_called.append(expected_disconnect)
async def on_connect() -> None:
nonlocal on_connect_called
on_connect_called.append(True)
async def on_connect_fail(connect_exception: Exception) -> None:
nonlocal on_connect_called
on_connect_fail_called.append(connect_exception)
rl = ReconnectLogic(
client=cli,
on_disconnect=on_disconnect,
on_connect=on_connect,
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
on_connect_error=on_connect_fail,
)
assert cli.log_name == "mydevice @ 1.2.3.4"
caplog.clear()
with patch.object(cli, "start_connection", side_effect=APIConnectionError):
await rl.start()
await asyncio.sleep(0)
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 0
assert len(on_connect_fail_called) == 1
assert isinstance(on_connect_fail_called[-1], APIConnectionError)
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert "connect to ESPHome API for mydevice @ 1.2.3.4" in caplog.text
for record in caplog.records:
if "connect to ESPHome API for mydevice @ 1.2.3.4" in record.message:
assert record.levelno == logging.WARNING
caplog.clear()
# Next retry should run at debug level
with patch.object(cli, "start_connection", side_effect=APIConnectionError):
# Should now retry
assert rl._connect_timer is not None
rl._connect_timer._run()
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 0
assert len(on_connect_fail_called) == 2
assert isinstance(on_connect_fail_called[-1], APIConnectionError)
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert "connect to ESPHome API for mydevice @ 1.2.3.4" in caplog.text
for record in caplog.records:
if "connect to ESPHome API for mydevice @ 1.2.3.4" in record.message:
assert record.levelno == logging.DEBUG
caplog.clear()
with patch.object(cli, "start_connection"), patch.object(cli, "finish_connection"):
# Should now retry
assert rl._connect_timer is not None
rl._connect_timer._run()
await asyncio.sleep(0)
await asyncio.sleep(0)
assert "connect to ESPHome API for mydevice @ 1.2.3.4" not in caplog.text
assert len(on_disconnect_called) == 0
assert len(on_connect_called) == 1
assert len(on_connect_fail_called) == 2
assert rl._connection_state is ReconnectLogicState.READY
original_when = rl._connect_timer.when()
# Ensure starting the connection logic again does not trigger a new connection
await rl.start()
# Verify no new timer is started
assert rl._connect_timer.when() == original_when
await rl.stop()
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
DNS_POINTER = DNSPointer(
"_esphomelib._tcp.local.",
_TYPE_PTR,
_CLASS_IN,
1000,
"mydevice._esphomelib._tcp.local.",
)
@pytest.mark.parametrize(
("record", "should_trigger_zeroconf", "expected_state_after_trigger", "log_text"),
(
(
DNS_POINTER,
True,
ReconnectLogicState.READY,
"received mDNS record",
),
(
DNSPointer(
"_esphomelib._tcp.local.",
_TYPE_PTR,
_CLASS_IN,
1000,
"wrong_name._esphomelib._tcp.local.",
),
False,
ReconnectLogicState.CONNECTING,
"",
),
(
DNSAddress(
"mydevice.local.",
_TYPE_A,
_CLASS_IN,
1000,
ip_address("1.2.3.4").packed,
),
True,
ReconnectLogicState.READY,
"received mDNS record",
),
),
)
@pytest.mark.asyncio
async def test_reconnect_zeroconf(
patchable_api_client: APIClient,
caplog: pytest.LogCaptureFixture,
record: DNSRecord,
should_trigger_zeroconf: bool,
expected_state_after_trigger: ReconnectLogicState,
log_text: str,
) -> None:
"""Test that reconnect logic retry."""
cli = patchable_api_client
mock_zeroconf = MagicMock(spec=Zeroconf)
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
zeroconf_instance=mock_zeroconf,
name="mydevice",
on_connect_error=AsyncMock(),
)
assert cli.log_name == "mydevice @ 1.2.3.4"
with patch.object(
cli, "start_connection", side_effect=quick_connect_fail
) as mock_start_connection:
await rl.start()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
with patch.object(
cli, "start_connection", side_effect=slow_connect_fail
) as mock_start_connection:
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert rl._accept_zeroconf_records is True
assert not rl._is_stopped
assert rl._connect_timer is not None
rl._connect_timer._run()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
assert rl._connection_state is ReconnectLogicState.CONNECTING
assert rl._accept_zeroconf_records is True
assert not rl._is_stopped
caplog.clear()
with patch.object(cli, "start_connection") as mock_start_connection, patch.object(
cli, "finish_connection"
):
assert rl._zc_listening is True
rl.async_update_records(
mock_zeroconf, current_time_millis(), [RecordUpdate(record, None)]
)
assert (
"Triggering connect because of received mDNS record" in caplog.text
) is should_trigger_zeroconf
assert rl._accept_zeroconf_records is not should_trigger_zeroconf
assert rl._zc_listening is True # should change after one iteration of the loop
await asyncio.sleep(0)
assert rl._zc_listening is not should_trigger_zeroconf
# The reconnect is scheduled to run in the next loop iteration
await asyncio.sleep(0)
assert mock_start_connection.call_count == int(should_trigger_zeroconf)
assert log_text in caplog.text
assert rl._connection_state is expected_state_after_trigger
await rl.stop()
assert rl._is_stopped is True
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_reconnect_zeroconf_not_while_handshaking(
patchable_api_client: APIClient,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that reconnect logic retry will not trigger a zeroconf reconnect while handshaking."""
cli = patchable_api_client
mock_zeroconf = MagicMock(spec=Zeroconf)
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
zeroconf_instance=mock_zeroconf,
name="mydevice",
on_connect_error=AsyncMock(),
)
assert cli.log_name == "mydevice @ 1.2.3.4"
with patch.object(
cli, "start_connection", side_effect=quick_connect_fail
) as mock_start_connection:
await rl.start()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
with patch.object(cli, "start_connection") as mock_start_connection, patch.object(
cli, "finish_connection", side_effect=slow_connect_fail
) as mock_finish_connection:
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert rl._accept_zeroconf_records is True
assert not rl._is_stopped
assert rl._connect_timer is not None
rl._connect_timer._run()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
assert mock_finish_connection.call_count == 1
assert rl._connection_state is ReconnectLogicState.HANDSHAKING
assert rl._accept_zeroconf_records is False
assert not rl._is_stopped
rl.async_update_records(
mock_zeroconf, current_time_millis(), [RecordUpdate(DNS_POINTER, None)]
)
assert (
"Triggering connect because of received mDNS record" in caplog.text
) is False
rl._cancel_connect("forced cancel in test")
await rl.stop()
assert rl._is_stopped is True
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_connect_task_not_cancelled_while_handshaking(
patchable_api_client: APIClient,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that reconnect logic will not cancel an in progress handshake."""
cli = patchable_api_client
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
name="mydevice",
on_connect_error=AsyncMock(),
)
assert cli.log_name == "mydevice @ 1.2.3.4"
with patch.object(
cli, "start_connection", side_effect=quick_connect_fail
) as mock_start_connection:
await rl.start()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
with patch.object(cli, "start_connection") as mock_start_connection, patch.object(
cli, "finish_connection", side_effect=slow_connect_fail
) as mock_finish_connection:
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
assert rl._accept_zeroconf_records is True
assert not rl._is_stopped
assert rl._connect_timer is not None
rl._connect_timer._run()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
assert mock_finish_connection.call_count == 1
assert rl._connection_state is ReconnectLogicState.HANDSHAKING
assert rl._accept_zeroconf_records is False
assert not rl._is_stopped
caplog.clear()
# This can likely never happen in practice, but we should handle it
# in the event there is a race as the consequence is that we could
# disconnect a working connection.
rl._call_connect_once()
assert (
"Not cancelling existing connect task as its already ReconnectLogicState.HANDSHAKING"
in caplog.text
)
rl._cancel_connect("forced cancel in test")
await rl.stop()
assert rl._is_stopped is True
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_connect_aborts_if_stopped(
patchable_api_client: APIClient,
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that reconnect logic will abort connecting if stopped."""
cli = patchable_api_client
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
name="mydevice",
on_connect_error=AsyncMock(),
)
assert cli.log_name == "mydevice @ 1.2.3.4"
with patch.object(
cli, "start_connection", side_effect=quick_connect_fail
) as mock_start_connection:
await rl.start()
await asyncio.sleep(0)
assert mock_start_connection.call_count == 1
with patch.object(cli, "start_connection") as mock_start_connection:
timer = rl._connect_timer
assert timer is not None
await rl.stop()
assert rl._is_stopped is True
rl._call_connect_once()
await asyncio.sleep(0)
await asyncio.sleep(0)
# We should never try to connect again
# once we are stopped
assert mock_start_connection.call_count == 0
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_reconnect_logic_stop_callback(patchable_api_client: APIClient):
"""Test that the stop_callback stops the ReconnectLogic."""
cli = patchable_api_client
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
)
await rl.start()
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
await asyncio.sleep(0)
assert rl._connection_state is ReconnectLogicState.CONNECTING
assert rl._is_stopped is False
rl.stop_callback()
# Wait for cancellation to propagate
for _ in range(4):
await asyncio.sleep(0)
assert rl._is_stopped is True
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_reconnect_logic_stop_callback_waits_for_handshake(
patchable_api_client: APIClient,
):
"""Test that the stop_callback waits for a handshake."""
cli = patchable_api_client
rl = ReconnectLogic(
client=cli,
on_disconnect=AsyncMock(),
on_connect=AsyncMock(),
zeroconf_instance=get_mock_zeroconf(),
name="mydevice",
)
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
with patch.object(cli, "start_connection"), patch.object(
cli, "finish_connection", side_effect=slow_connect_fail
):
await rl.start()
for _ in range(3):
await asyncio.sleep(0)
assert rl._connection_state is ReconnectLogicState.HANDSHAKING
assert rl._is_stopped is False
rl.stop_callback()
# Wait for cancellation to propagate
for _ in range(4):
await asyncio.sleep(0)
assert rl._is_stopped is False
assert rl._connection_state is ReconnectLogicState.HANDSHAKING
rl._cancel_connect("forced cancel in test")
# Wait for cancellation to propagate
for _ in range(4):
await asyncio.sleep(0)
assert rl._is_stopped is True
assert rl._connection_state is ReconnectLogicState.DISCONNECTED
@pytest.mark.asyncio
async def test_handling_unexpected_disconnect(
event_loop: asyncio.AbstractEventLoop, aiohappyeyeballs_start_connection
):
"""Test the disconnect callback fires with expected_disconnect=False."""
loop = asyncio.get_event_loop()
protocol: APIPlaintextFrameHelper | None = None
transport = MagicMock()
connected = asyncio.Event()
class PatchableAPIClient(APIClient):
pass
async_zeroconf = get_mock_async_zeroconf()
cli = PatchableAPIClient(
address="1.2.3.4",
port=6052,
password=None,
noise_psk=None,
expected_name="fake",
zeroconf_instance=async_zeroconf.zeroconf,
)
connected = asyncio.Event()
on_disconnect_calls = []
async def on_disconnect(expected_disconnect: bool) -> None:
on_disconnect_calls.append(expected_disconnect)
async def on_connect() -> None:
connected.set()
logic = ReconnectLogic(
client=cli,
on_connect=on_connect,
on_disconnect=on_disconnect,
zeroconf_instance=async_zeroconf,
name="fake",
)
with patch.object(
loop,
"create_connection",
side_effect=partial(_create_mock_transport_protocol, transport, connected),
):
await logic.start()
await connected.wait()
protocol = cli._connection._frame_helper
send_plaintext_hello(protocol)
send_plaintext_connect_response(protocol, False)
await asyncio.sleep(0)
await asyncio.sleep(0)
assert cli._connection.is_connected is True
await asyncio.sleep(0)
with patch.object(
loop,
"create_connection",
side_effect=partial(_create_mock_transport_protocol, transport, connected),
) as mock_create_connection:
protocol.eof_received()
# Wait for the task to run
await asyncio.sleep(0)
# Ensure we try to reconnect immediately
# since its an unexpected disconnect
assert mock_create_connection.call_count == 0
assert len(on_disconnect_calls) == 1
expected_disconnect = on_disconnect_calls[-1]
assert expected_disconnect is False
await logic.stop()
@pytest.mark.asyncio
async def test_backoff_on_encryption_error(
event_loop: asyncio.AbstractEventLoop,
caplog: pytest.LogCaptureFixture,
aiohappyeyeballs_start_connection,
) -> None:
"""Test we backoff on encryption error."""
loop = asyncio.get_event_loop()
protocol: APIPlaintextFrameHelper | None = None
transport = MagicMock()
connected = asyncio.Event()
class PatchableAPIClient(APIClient):
pass
async_zeroconf = get_mock_async_zeroconf()
cli = PatchableAPIClient(
address="1.2.3.4",
port=6052,
password=None,
noise_psk="",
expected_name="fake",
zeroconf_instance=async_zeroconf.zeroconf,
)
connected = asyncio.Event()
on_disconnect_calls = []
async def on_disconnect(expected_disconnect: bool) -> None:
on_disconnect_calls.append(expected_disconnect)
async def on_connect() -> None:
connected.set()
logic = ReconnectLogic(
client=cli,
on_connect=on_connect,
on_disconnect=on_disconnect,
zeroconf_instance=async_zeroconf,
name="fake",
)
with patch.object(
loop,
"create_connection",
side_effect=partial(_create_mock_transport_protocol, transport, connected),
):
await logic.start()
await connected.wait()
protocol = cli._connection._frame_helper
mock_data_received(protocol, b"\x01\x00\x00")
assert cli._connection.is_connected is False
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(on_disconnect_calls) == 0
assert "Scheduling new connect attempt in 60.00 seconds" in caplog.text
assert "Connection requires encryption (RequiresEncryptionAPIError)" in caplog.text
now = loop.time()
assert logic._connect_timer.when() - now == pytest.approx(60, 1)
assert logic._tries == MAXIMUM_BACKOFF_TRIES
await logic.stop()