From 00955d449ae24a7c4cf66c260bbde887d542a9b5 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 10 Oct 2026 13:06:14 -1000 Subject: [PATCH] Cover the UDP client paths in the espota2 unit tests --- tests/unit_tests/test_espota2.py | 168 +++++++++++++++++++++++++ tests/unit_tests/test_espota2_noise.py | 54 +++++++- 2 files changed, 221 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py index 6373c00ff8d..3efe0f91a75 100644 --- a/tests/unit_tests/test_espota2.py +++ b/tests/unit_tests/test_espota2.py @@ -1751,3 +1751,171 @@ def test_udp_channel_reports_socket_errors() -> None: ): channel.sendall(bytes(espota2.UDP_MAX_PAYLOAD * 5)) channel.close() + + +def _udp_handshake() -> list[bytes]: + """TCP recv responses up to auth, from a device that offers UDP.""" + return [ + bytes([espota2.RESPONSE_OK]), + bytes([espota2.OTA_VERSION_2_0]), + bytes([espota2.RESPONSE_FEATURE_FLAGS]), + bytes([espota2.SERVER_FEATURE_SUPPORTS_UDP]), + bytes([espota2.RESPONSE_AUTH_OK]), + ] + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_over_udp(mock_socket: Mock) -> None: + """The data phase moves to the channel, which acks chunks itself.""" + mock_socket.recv.side_effect = _udp_handshake() + udp = Mock() + udp.pending.return_value = False + udp.recv.side_effect = [ + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + bytes([espota2.RESPONSE_RECEIVE_OK]), + bytes([espota2.RESPONSE_UPDATE_END_OK]), + ] + with patch.object(espota2, "_start_udp", return_value=udp) as start: + espota2.perform_ota( + mock_socket, None, io.BytesIO(bytes(20000)), "test.bin", prefer_udp=True + ) + assert start.call_args.args[2] is True # an earlier failure means lossy + assert udp.recv.call_count == 4 # no chunk acks + udp.flush.assert_called_once() + udp.close.assert_called_once() + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_over_udp_device_error(mock_socket: Mock) -> None: + """A device error mid upload stops the upload and is not retried.""" + mock_socket.recv.side_effect = _udp_handshake() + mock_socket.getsockopt.return_value = b"" # no retransmit counter + udp = Mock() + udp.pending.return_value = True + udp.recv.side_effect = [ + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]), + ] + with ( + patch.object(espota2, "_start_udp", return_value=udp), + pytest.raises(espota2.OTAError, match="Writing OTA data") as exc, + ): + espota2.perform_ota(mock_socket, None, io.BytesIO(bytes(20000)), "test.bin") + assert not isinstance(exc.value, espota2.OTANetworkError) + udp.close.assert_called() + + +@pytest.mark.usefixtures("mock_time") +def test_perform_ota_udp_declined_stays_on_tcp(mock_socket: Mock) -> None: + """Without a channel the upload continues on TCP with chunk acks.""" + mock_socket.recv.side_effect = [ + *_udp_handshake(), + bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]), + bytes([espota2.RESPONSE_BIN_MD5_OK]), + *_UPLOAD_TAIL, + ] + with patch.object(espota2, "_start_udp", return_value=None): + espota2.perform_ota( + mock_socket, None, io.BytesIO(b"x" * 100), "test.bin", prefer_udp=True + ) + assert mock_socket.sendall.call_args_list[-1] == call(bytes([espota2.RESPONSE_OK])) + + +def test_udp_channel_records_socket_errors() -> None: + """Send and receive errors are recorded like loss.""" + sock = Mock() + sock.send.side_effect = BlockingIOError(35, "full") + sock.recv.side_effect = ConnectionRefusedError(61, "refused") + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel._send(espota2.UDP_MSG_PING, 0, b"") + assert isinstance(channel._last_error, BlockingIOError) + with patch("select.select", side_effect=[([sock], [], []), ([], [], [])]): + channel._receive(0.0) + assert isinstance(channel._last_error, ConnectionRefusedError) + + +def test_start_udp_probe_succeeds(mock_socket: Mock) -> None: + """A device that answers the probe gets the data phase over UDP.""" + with _LossyUdpDevice(b"\x05\x06\x07\x09", 0.0) as device: + mock_socket.family = socket.AF_INET + mock_socket.getpeername.return_value = device.sock.getsockname() + mock_socket.recv.side_effect = [b"\x05", b"\x06\x07\x09"] + channel = espota2._start_udp(mock_socket, mock_socket, True) + assert channel is not None + channel.close() + assert mock_socket.sendall.call_args_list == [call(espota2.UDP_REQUEST)] + + +def test_udp_channel_pings_when_all_held() -> None: + """With every message in flight already held, it asks for responses instead, + at most once per resend timeout.""" + device = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + device.bind(("127.0.0.1", 0)) + device.settimeout(1.0) + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.connect(device.getsockname()) + channel = espota2.UdpChannel(sock, b"\x01\x01\x01\x01") + channel.sendall(b"x") + channel._queue[0].held = True + with patch.object(channel, "_receive"): + channel._pump() + channel._pump() + assert device.recv(2048)[0] == espota2.UDP_MSG_PING + device.settimeout(0.1) + with pytest.raises(TimeoutError): + device.recv(2048) + channel.close() + device.close() + + +def test_udp_channel_ignores_foreign_datagrams() -> None: + """Short, non-ACK or wrong-token datagrams change nothing.""" + channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01") + for data in (b"\x40", bytes(espota2.UDP_ACK_HEADER_SIZE), b"\x40" + bytes(11)): + channel._handle_ack(data) + assert not channel.pending() + # An ACK before the head was ever sent is no duplicate + channel.sendall(b"x") + ack = bytes([espota2.UDP_MSG_ACK]) + b"\x01\x01\x01\x01" + bytes(6) + b"\x40" + channel._handle_ack(ack) + channel._handle_ack(ack) + assert not channel._queue[0].resend + + +def test_udp_channel_logs_each_socket_error_once( + caplog: pytest.LogCaptureFixture, +) -> None: + """A repeated error is recorded but logged only when it changes.""" + channel = espota2.UdpChannel(Mock(), b"\x01\x01\x01\x01") + with caplog.at_level(logging.DEBUG, "esphome.espota2"): + channel._socket_error(ConnectionRefusedError(61, "refused")) + channel._socket_error(ConnectionRefusedError(61, "refused again")) + assert caplog.text.count("UDP socket error") == 1 + assert "again" in str(channel._last_error) + + +def test_tcp_retransmits_unsupported_option() -> None: + """An OS that rejects the option reports nothing.""" + sock = Mock() + sock.getsockopt.side_effect = OSError("not supported") + with patch("sys.platform", "linux"): + assert espota2._tcp_retransmits(sock) is None + + +@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip") +def test_run_ota_impl_handshake_failure_prefers_udp( + firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock +) -> None: + """A handshake that dies after reaching the device moves the retry to UDP.""" + mock_perform_ota.side_effect = [ + espota2.OTAHandshakeNetworkError("noise handshake: timed out"), + None, + ] + result_code, _ = espota2.run_ota_impl_("test.local", 3232, None, str(firmware_file)) + assert result_code == 0 + assert [c.kwargs["prefer_udp"] for c in mock_perform_ota.call_args_list] == [ + False, + True, + ] diff --git a/tests/unit_tests/test_espota2_noise.py b/tests/unit_tests/test_espota2_noise.py index 60f6b28e076..5aba009848b 100644 --- a/tests/unit_tests/test_espota2_noise.py +++ b/tests/unit_tests/test_espota2_noise.py @@ -70,8 +70,10 @@ class FakeEncryptedDevice(threading.Thread): prologue_features_override: int | None = None, connections: int = 1, drop_handshakes: int = 0, + udp: bool = False, ) -> None: super().__init__(daemon=True) + self.udp = udp # offer UDP; the data phase then has no chunk acks self.connections = connections self.drop_handshakes = drop_handshakes # hang up mid-handshake this many times self.psk = psk @@ -116,6 +118,8 @@ class FakeEncryptedDevice(threading.Thread): sock.sendall(bytes([espota2.RESPONSE_ERROR_ENCRYPTION_REQUIRED])) return server_flags = espota2.SERVER_FEATURE_SUPPORTS_NOISE if self.offer_noise else 0 + if self.udp: + server_flags |= espota2.SERVER_FEATURE_SUPPORTS_UDP sock.sendall(bytes([espota2.RESPONSE_FEATURE_FLAGS, server_flags])) if not (noise_negotiated and self.offer_noise): # A device that does not require encryption continues in @@ -200,7 +204,7 @@ class FakeEncryptedDevice(threading.Thread): acked = 0 while len(received) < size: received += recv_data(size - len(received)) - if self.version >= espota2.OTA_VERSION_2_0: + if self.version >= espota2.OTA_VERSION_2_0 and not self.udp: while acked + espota2.UPLOAD_BLOCK_SIZE <= len(received) or ( len(received) == size and acked < size ): @@ -598,3 +602,51 @@ def test_bare_block_refuses_a_device_that_cannot_encrypt( assert device.received != b"firmware" assert any("refusing to send the image" in r.message for r in caplog.records) assert not any("Retrying in plaintext" in r.message for r in caplog.records) + + +class _TcpAsUdp: + """Stands in for a UdpChannel, carrying its byte stream over the TCP socket.""" + + def __init__(self, sock: socket.socket) -> None: + self._sock = sock + self.flushed = False + self.closed = False + + def settimeout(self, timeout: float | None) -> None: + self._sock.settimeout(timeout) + + def setsockopt(self, level: int, optname: int, value: int) -> None: + self._sock.setsockopt(level, optname, value) + + def sendall(self, data: bytes) -> None: + self._sock.sendall(data) + + def recv(self, amount: int) -> bytes: + return self._sock.recv(amount) + + def pending(self) -> bool: + return False + + def flush(self, timeout: float) -> None: + self.flushed = True + + def close(self) -> None: + self.closed = True + + +def test_encrypted_upload_over_udp() -> None: + """After the handshake, Noise frames continue on the UDP channel.""" + pytest.importorskip("aioesphomeapi.noise") + firmware = bytes(range(256)) * 80 + device = FakeEncryptedDevice(udp=True) + channels: list[_TcpAsUdp] = [] + + def start_udp(sock: Any, tcp_sock: socket.socket, lossy: bool) -> _TcpAsUdp: + channels.append(_TcpAsUdp(tcp_sock)) + return channels[0] + + with patch("time.sleep"), patch.object(espota2, "_start_udp", start_udp): + _upload(device, firmware, PSK) + device.join_and_check() + assert device.received == firmware + assert channels[0].flushed and channels[0].closed