[ota] Retry uploads that fail from network errors (#18332)

esphome/espota2.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+import contextlib
 import gzip
 import hashlib
 import io
@@ -8,7 +9,6 @@ import logging
 from pathlib import Path
 import secrets
 import socket
-import sys
 import time
 from typing import Any
 
@@ -76,6 +76,14 @@ _SUPPORTED_OTA_TYPES: frozenset[int] = frozenset(
 UPLOAD_BLOCK_SIZE = 8192
 UPLOAD_BUFFER_SIZE = UPLOAD_BLOCK_SIZE * 8
 
+# Flaky Wi-Fi links often drop the first OTA attempt, and the device may need time
+# to clean up a half-open connection (its handshake watchdog runs at 20s) before it
+# accepts a new one, so wait between attempts instead of failing the upload outright.
+# Every resolved address is tried once, and this many extra attempts are shared
+# across the addresses on top of that.
+EXTRA_UPLOAD_ATTEMPTS = 2
+UPLOAD_RETRY_DELAY = 5.0
+
 _LOGGER = logging.getLogger(__name__)
 
 # Authentication method lookup table: response -> (hash_func, nonce_size, name)
@@ -171,6 +179,23 @@ class OTAError(EsphomeError):
     pass
 
 
+class OTANetworkError(OTAError):
+    """Network-level OTA failure (timeout, reset, closed connection); retrying may succeed."""
+
+
+def _committed_error(err: OTANetworkError) -> OTAError:
+    """Wrap a network failure that happened once the device had the full image.
+
+    Past that point the device commits and reboots on its own, so the failure
+    must not be retried; a re-upload could flash a device that already updated.
+    """
+    return OTAError(
+        f"{err} (the device may have already committed the update and "
+        f"be rebooting; check whether it comes back with the new "
+        f"firmware before uploading again)"
+    )
+
+
 def recv_decode(
     sock: socket.socket, amount: int, decode: bool = True
 ) -> bytes | list[int]:
@@ -209,19 +234,22 @@ def receive_exactly(
     try:
         data += recv_decode(sock, 1, decode=decode)  # type: ignore[operator]
     except OSError as err:
-        raise OTAError(f"receiving {msg} response: {err}") from err
+        raise OTANetworkError(f"receiving {msg} response: {err}") from err
 
     try:
         check_error(data, expect)
     except OTAError as err:
         sock.close()
-        raise OTAError(f"receiving {msg}: {err}") from err
+        # type(err) preserves OTANetworkError vs OTAError so callers can tell
+        # retryable network failures from device-reported errors; subclasses
+        # must accept a single message argument
+        raise type(err)(f"receiving {msg}: {err}") from err
 
     while len(data) < amount:
         try:
             data += recv_decode(sock, amount - len(data), decode=decode)  # type: ignore[operator]
         except OSError as err:
-            raise OTAError(f"receiving {msg}: {err}") from err
+            raise OTANetworkError(f"receiving {msg}: {err}") from err
     return data
 
 
@@ -237,7 +265,7 @@ def check_error(data: list[int] | bytes, expect: int | list[int] | None) -> None
     # accept-any-response reads (e.g. feature negotiation, auth nonces) would be
     # silently passed through and surface later as cryptic decode/timeout failures.
     if not data:
-        raise OTAError(
+        raise OTANetworkError(
             "Device closed connection without responding. "
             "This may indicate the device ran out of memory, "
             "a network issue, or the connection was interrupted."
@@ -274,7 +302,7 @@ def send_check(
 
         sock.sendall(data)
     except OSError as err:
-        raise OTAError(f"sending {msg}: {err}") from err
+        raise OTANetworkError(f"sending {msg}: {err}") from err
 
 
 def perform_ota(
@@ -306,7 +334,7 @@ def perform_ota(
     send_check(sock, MAGIC_BYTES, "magic bytes")
 
     _, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
-    _LOGGER.debug("Device support OTA version: %s", version)
+    _LOGGER.info("Connection established; device supports OTA version %s", version)
     supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0)
     if version not in supported_versions:
         raise OTAError(
@@ -417,6 +445,8 @@ def perform_ota(
         hash_func, nonce_size, hash_name = _AUTH_METHODS[auth]
         perform_auth(sock, password, hash_func, nonce_size, hash_name)
 
+    _LOGGER.info("Handshake complete")
+
     # Timeout must match device-side OTA_SOCKET_TIMEOUT_DATA to prevent premature failures
     sock.settimeout(90.0)
 
@@ -449,21 +479,43 @@ def perform_ota(
 
     offset = 0
     progress = ProgressBar("Uploading")
-    while True:
-        chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE]
-        if not chunk:
-            break
-        offset += len(chunk)
+    try:
+        while True:
+            chunk = upload_contents[offset : offset + UPLOAD_BLOCK_SIZE]
+            if not chunk:
+                break
+            offset += len(chunk)
+
+            try:
+                sock.sendall(chunk)
+            except OSError as err:
+                # A send failure can hide an error byte the device reported
+                # just before dropping the connection; surface that as the
+                # real, non-retryable cause when it is available
+                try:
+                    sock.settimeout(1.0)
+                    check_error(recv_decode(sock, 1), None)
+                except (OSError, OTANetworkError) as probe_err:
+                    _LOGGER.debug(
+                        "No device error behind the send failure: %s", probe_err
+                    )
+                raise OTANetworkError(f"sending data: {err}") from err
 
-        try:
-            sock.sendall(chunk)
             if version >= OTA_VERSION_2_0:
-                receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
-        except OSError as err:
-            sys.stderr.write("\n")
-            raise OTAError(f"sending data: {err}") from err
+                try:
+                    receive_exactly(sock, 1, "chunk result", RESPONSE_CHUNK_OK)
+                except OTANetworkError as err:
+                    if offset < upload_size:
+                        raise
+                    # The device already had the complete image when this ack
+                    # was lost, so it may be committing; do not retry
+                    raise _committed_error(err) from err
 
-        progress.update(offset / upload_size)
+            progress.update(offset / upload_size)
+    except OTAError:
+        # Terminate the progress bar line before the error is logged
+        progress.done()
+        raise
     progress.done()
 
     # Enable nodelay for last checks
@@ -472,11 +524,25 @@ def perform_ota(
 
     _LOGGER.info("Upload took %.2f seconds, waiting for result...", duration)
 
-    receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK)
-    receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK)
-    send_check(sock, RESPONSE_OK, "end acknowledgement")
+    # Once the device has the complete image it commits the update and
+    # reboots on its own; the exact commit point is not observable from
+    # here, so treat everything past the data phase as non-retryable. A
+    # re-upload could flash a device that already updated successfully.
+    try:
+        receive_exactly(sock, 1, "update receive result", RESPONSE_RECEIVE_OK)
+        receive_exactly(sock, 1, "update end result", RESPONSE_UPDATE_END_OK)
+    except OTANetworkError as err:
+        raise _committed_error(err) from err
 
-    _LOGGER.info("OTA successful")
+    try:
+        send_check(sock, RESPONSE_OK, "end acknowledgement")
+    except OTANetworkError as err:
+        # The device treats a missing end acknowledgement as non-fatal and is
+        # already rebooting into the new firmware, so the update succeeded
+        _LOGGER.warning("Failed sending end acknowledgement: %s", err)
+        _LOGGER.info("OTA successful (end acknowledgement not delivered)")
+    else:
+        _LOGGER.info("OTA successful")
 
     # Do not connect logs until it is fully on
     time.sleep(1)
@@ -510,8 +576,33 @@ def run_ota_impl_(
         )
         raise OTAError(err) from err
 
-    for r in res:
-        af, socktype, _, _, sa = r
+    if not res:
+        _LOGGER.error("No addresses to connect to for %s", remote_host)
+        return 1, None
+
+    # Every address is tried at least once and EXTRA_UPLOAD_ATTEMPTS retries
+    # are shared across the addresses, cycling through them. Wait before an
+    # attempt when the previous one actually reached the device, or when
+    # revisiting an address, so a flaky link can recover and the device can
+    # clean up a half-open connection (its handshake watchdog runs at 20s);
+    # moving on to the next address family stays immediate. Known limitation:
+    # a silent mid-transfer drop with no reset can wedge the device until its
+    # 90s data timeout, which outlasts this budget; the retries target the
+    # common failures where the device resets or closes the link promptly.
+    total_attempts = len(res) + EXTRA_UPLOAD_ATTEMPTS
+    last_error = ""
+    reached_device = False
+    for attempt in range(total_attempts):
+        af, socktype, _, _, sa = res[attempt % len(res)]
+        if reached_device or attempt >= len(res):
+            _LOGGER.info(
+                "Retrying in %.0f seconds (attempt %d of %d)...",
+                UPLOAD_RETRY_DELAY,
+                attempt + 1,
+                total_attempts,
+            )
+            time.sleep(UPLOAD_RETRY_DELAY)
+        reached_device = False
         _LOGGER.info("Connecting to %s port %s...", sa[0], sa[1])
         sock = socket.socket(af, socktype)
         sock.settimeout(20.0)
@@ -519,23 +610,30 @@ def run_ota_impl_(
             sock.connect(sa)
         except OSError as err:
             sock.close()
-            _LOGGER.error("Connecting to %s port %s failed: %s", sa[0], sa[1], err)
+            _LOGGER.warning("Connecting to %s port %s failed: %s", sa[0], sa[1], err)
+            last_error = f"connecting to {sa[0]} failed: {err}"
             continue
 
         _LOGGER.info("Connected to %s", sa[0])
-        with Path(filename).open("rb") as file_handle:
+        reached_device = True
+        with contextlib.closing(sock), Path(filename).open("rb") as file_handle:
             try:
                 perform_ota(sock, password, file_handle, filename, ota_type)
+            except OTANetworkError as err:
+                # Transient network failure; retry
+                last_error = str(err)
+                _LOGGER.warning("%s", last_error)
+                continue
             except OTAError as err:
+                # Device-reported error (wrong password, wrong flash size, ...);
+                # retrying cannot succeed, so fail immediately
                 _LOGGER.error(str(err))
                 return 1, None
-            finally:
-                sock.close()
 
         # Successfully uploaded to sa[0]
         return 0, sa[0]
 
-    _LOGGER.error("Connection failed.")
+    _LOGGER.error("Upload failed after %d attempts: %s", total_attempts, last_error)
     return 1, None
 
 

tests/unit_tests/test_espota2.py

@@ -44,13 +44,17 @@ def mock_file() -> io.BytesIO:
 
 
 @pytest.fixture
-def mock_time() -> Generator[None]:
+def mock_sleep() -> Generator[Mock]:
+    """Mock time.sleep so delays don't slow down tests."""
+    with patch("time.sleep") as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_time(mock_sleep: Mock) -> Generator[None]:
     """Mock time-related functions for consistent testing."""
     # Provide enough values for multiple calls (tests may call perform_ota multiple times)
-    with (
-        patch("time.sleep"),
-        patch("time.perf_counter", side_effect=[0, 1, 0, 1, 0, 1]),
-    ):
+    with patch("time.perf_counter", side_effect=[0, 1, 0, 1, 0, 1]):
         yield
 
 
@@ -79,6 +83,28 @@ def mock_resolve_ip() -> Generator[Mock]:
         yield mock
 
 
+DUAL_STACK_SA6 = ("2001:db8::1", 3232, 0, 0)
+DUAL_STACK_SA4 = ("192.168.1.100", 3232)
+
+
+@pytest.fixture
+def mock_resolve_ip_dual(mock_resolve_ip: Mock) -> Mock:
+    """Make resolve_ip_address return an IPv6 and an IPv4 address."""
+    mock_resolve_ip.return_value = [
+        (socket.AF_INET6, socket.SOCK_STREAM, 0, "", DUAL_STACK_SA6),
+        (socket.AF_INET, socket.SOCK_STREAM, 0, "", DUAL_STACK_SA4),
+    ]
+    return mock_resolve_ip
+
+
+@pytest.fixture
+def firmware_file(tmp_path: Path) -> Path:
+    """Create a firmware file on disk for run_ota_impl_ tests."""
+    firmware = tmp_path / "firmware.bin"
+    firmware.write_bytes(b"firmware content")
+    return firmware
+
+
 @pytest.fixture
 def mock_perform_ota() -> Generator[Mock]:
     """Mock perform_ota function for testing."""
@@ -137,9 +163,11 @@ def test_receive_exactly_with_error_response(mock_socket: Mock) -> None:
 
     with pytest.raises(
         espota2.OTAError, match="receiving auth:.*Authentication invalid"
-    ):
+    ) as exc_info:
         espota2.receive_exactly(mock_socket, 1, "auth", [espota2.RESPONSE_OK])
 
+    # Device-reported errors must stay plain OTAError, not the retryable kind
+    assert not isinstance(exc_info.value, espota2.OTANetworkError)
     mock_socket.close.assert_called_once()
 
 
@@ -147,10 +175,30 @@ def test_receive_exactly_socket_error(mock_socket: Mock) -> None:
     """Test receive_exactly handles socket errors."""
     mock_socket.recv.side_effect = OSError("Connection reset")
 
-    with pytest.raises(espota2.OTAError, match="receiving test response"):
+    with pytest.raises(espota2.OTANetworkError, match="receiving test response"):
         espota2.receive_exactly(mock_socket, 1, "test", espota2.RESPONSE_OK)
 
 
+def test_receive_exactly_mid_read_socket_error(mock_socket: Mock) -> None:
+    """Test receive_exactly handles socket errors after the first byte."""
+    mock_socket.recv.side_effect = [b"\x00", OSError("Connection reset")]
+
+    with pytest.raises(espota2.OTANetworkError, match="receiving test:"):
+        espota2.receive_exactly(mock_socket, 3, "test", espota2.RESPONSE_OK)
+
+
+def test_receive_exactly_closed_connection_is_network_error(mock_socket: Mock) -> None:
+    """Test receive_exactly raises OTANetworkError when the device closes the connection."""
+    mock_socket.recv.return_value = b""
+
+    with pytest.raises(
+        espota2.OTANetworkError, match="Device closed connection without responding"
+    ):
+        espota2.receive_exactly(mock_socket, 1, "test", espota2.RESPONSE_OK)
+
+    mock_socket.close.assert_called_once()
+
+
 @pytest.mark.parametrize(
     ("error_code", "expected_msg"),
     [
@@ -227,15 +275,15 @@ def test_check_error_unexpected_response() -> None:
 
 
 def test_check_error_empty_data() -> None:
-    """Test check_error raises error when device closes connection without responding."""
+    """Test check_error raises the retryable OTANetworkError when the device closes the connection."""
     with pytest.raises(
-        espota2.OTAError, match="Device closed connection without responding"
+        espota2.OTANetworkError, match="Device closed connection without responding"
     ):
         espota2.check_error([], [espota2.RESPONSE_OK])
 
     # Also test with empty bytes
     with pytest.raises(
-        espota2.OTAError, match="Device closed connection without responding"
+        espota2.OTANetworkError, match="Device closed connection without responding"
     ):
         espota2.check_error(b"", [espota2.RESPONSE_OK])
 
@@ -530,6 +578,144 @@ def test_perform_ota_upload_error(mock_socket: Mock, mock_file: io.BytesIO) -> N
         espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
 
 
+def _no_auth_handshake(version: int) -> list[bytes]:
+    """Recv responses for a handshake without auth, up to the MD5 check."""
+    return [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([version]),  # Version number
+        bytes([espota2.RESPONSE_HEADER_OK]),  # Features response
+        bytes([espota2.RESPONSE_AUTH_OK]),  # No auth required
+        bytes([espota2.RESPONSE_UPDATE_PREPARE_OK]),  # Binary size OK
+        bytes([espota2.RESPONSE_BIN_MD5_OK]),  # MD5 checksum OK
+    ]
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_chunk_send_error(mock_socket: Mock, mock_file: io.BytesIO) -> None:
+    """Test OTA raises the retryable OTANetworkError when sending a chunk fails."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Probe for a pending error byte fails too
+    ]
+    # Sends before the data phase: magic bytes, features, binary size, MD5;
+    # fail on the fifth sendall, the first firmware chunk
+    mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")]
+
+    with pytest.raises(espota2.OTANetworkError, match="sending data:"):
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_chunk_send_error_surfaces_device_error(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a device error byte pending behind a send failure becomes the cause."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_1_0),
+        bytes([espota2.RESPONSE_ERROR_WRITING_FLASH]),  # Reason the device closed
+    ]
+    mock_socket.sendall.side_effect = [None] * 4 + [OSError("Broken pipe")]
+
+    with pytest.raises(
+        espota2.OTAError, match="Writing OTA data to flash memory failed"
+    ) as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # The device-reported error is not retryable
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_final_chunk_ack_failure_not_retryable(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a lost ack for the final chunk is not retried."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Ack for the only (final) chunk is lost
+    ]
+
+    with pytest.raises(espota2.OTAError, match="receiving chunk result") as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # The device already had the whole image, so it may be committing
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_intermediate_chunk_ack_failure_retryable(
+    mock_socket: Mock,
+) -> None:
+    """Test a lost ack for a non-final chunk stays retryable."""
+    # Two chunks: the firmware is larger than one upload block
+    big_file = io.BytesIO(b"x" * (espota2.UPLOAD_BLOCK_SIZE + 1))
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_2_0),
+        OSError("Connection reset"),  # Ack for the first of two chunks is lost
+    ]
+
+    with pytest.raises(espota2.OTANetworkError, match="receiving chunk result"):
+        espota2.perform_ota(mock_socket, None, big_file, "test.bin")
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_post_commit_failure_not_retryable(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a network failure after the device committed is a plain OTAError."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_1_0),
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Device received everything
+        OSError("Connection reset"),  # Connection lost waiting for end result
+    ]
+
+    with pytest.raises(espota2.OTAError, match="receiving update end result") as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # Must not be the retryable kind; the device is already rebooting
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_md5_mismatch_not_marked_committed(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test an MD5 mismatch keeps its own message and stays non-retryable."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_1_0),
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Device received everything
+        bytes([espota2.RESPONSE_ERROR_MD5_MISMATCH]),  # Device aborted the update
+    ]
+
+    with pytest.raises(espota2.OTAError, match="MD5 code mismatch") as exc:
+        espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    # The device aborted without committing, so the message must not claim
+    # the update may have been installed, and the error must not be retried
+    assert not isinstance(exc.value, espota2.OTANetworkError)
+    assert "committed" not in str(exc.value)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_end_ack_send_failure_is_success(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test a send failure on the final acknowledgement does not fail the OTA."""
+    mock_socket.recv.side_effect = [
+        *_no_auth_handshake(espota2.OTA_VERSION_1_0),
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Device received everything
+        bytes([espota2.RESPONSE_UPDATE_END_OK]),  # Update committed
+    ]
+    # Sends: magic bytes, features, binary size, MD5, one firmware chunk;
+    # fail on the sixth sendall, the end acknowledgement
+    mock_socket.sendall.side_effect = [None] * 5 + [OSError("Broken pipe")]
+
+    # Must not raise; the device treats a missing acknowledgement as non-fatal
+    espota2.perform_ota(mock_socket, None, mock_file, "test.bin")
+
+    assert mock_socket.sendall.call_count == 6
+
+
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
 def test_run_ota_impl_successful(
     mock_socket: Mock, tmp_path: Path, mock_perform_ota: Mock
@@ -564,21 +750,183 @@ def test_run_ota_impl_successful(
 
 
 @pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
-def test_run_ota_impl_connection_failed(mock_socket: Mock, tmp_path: Path) -> None:
-    """Test run_ota_impl_ when connection fails."""
+def test_run_ota_impl_connection_failed(
+    mock_socket: Mock, firmware_file: Path, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ retries when connection fails and eventually gives up."""
     mock_socket.connect.side_effect = OSError("Connection refused")
 
-    # Create a real firmware file
-    firmware_file = tmp_path / "firmware.bin"
-    firmware_file.write_bytes(b"firmware content")
-
     result_code, result_host = espota2.run_ota_impl_(
         "test.local", 3232, "password", str(firmware_file)
     )
 
     assert result_code == 1
     assert result_host is None
-    mock_socket.close.assert_called_once()
+    # A single address gets the whole attempt budget, with a delay before
+    # each revisit
+    assert mock_socket.connect.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_socket.close.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
+    mock_sleep.assert_called_with(espota2.UPLOAD_RETRY_DELAY)
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
+def test_run_ota_impl_connect_retry_succeeds(
+    mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ succeeds when a retry connects after a failed attempt."""
+    mock_socket.connect.side_effect = [OSError("Connection timed out"), None]
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 0
+    assert result_host == "192.168.1.100"
+    assert mock_socket.connect.call_count == 2
+    mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY)
+    mock_perform_ota.assert_called_once()
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
+def test_run_ota_impl_network_error_retry_succeeds(
+    mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ retries after a network error during the upload."""
+    mock_perform_ota.side_effect = [
+        espota2.OTANetworkError("receiving features: Device closed connection"),
+        None,
+    ]
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 0
+    assert result_host == "192.168.1.100"
+    assert mock_perform_ota.call_count == 2
+    mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY)
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
+def test_run_ota_impl_network_error_exhausts_attempts(
+    mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ gives up after all attempts hit network errors."""
+    mock_perform_ota.side_effect = espota2.OTANetworkError("sending data: broken pipe")
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 1
+    assert result_host is None
+    assert mock_perform_ota.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS + 1
+    assert mock_sleep.call_count == espota2.EXTRA_UPLOAD_ATTEMPTS
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
+def test_run_ota_impl_multiple_addresses_cycle(
+    mock_socket: Mock, firmware_file: Path, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ visits every address and cycles for the retries."""
+    mock_socket.connect.side_effect = OSError("No route to host")
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 1
+    assert result_host is None
+    # Each address is visited once, then the EXTRA_UPLOAD_ATTEMPTS spare
+    # attempts cycle back through them; the budget is shared, not per address
+    assert mock_socket.connect.call_args_list == [
+        call(DUAL_STACK_SA6),
+        call(DUAL_STACK_SA4),
+        call(DUAL_STACK_SA6),
+        call(DUAL_STACK_SA4),
+    ]
+    # No connect ever reached the device, so the delay only applies before
+    # the revisits
+    assert mock_sleep.call_count == 2
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
+def test_run_ota_impl_second_address_succeeds_without_delay(
+    mock_socket: Mock,
+    firmware_file: Path,
+    mock_perform_ota: Mock,
+    mock_sleep: Mock,
+) -> None:
+    """Test run_ota_impl_ falls through to the next address with no pause."""
+    mock_socket.connect.side_effect = [OSError("No route to host"), None]
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 0
+    assert result_host == "192.168.1.100"
+    mock_sleep.assert_not_called()
+    mock_perform_ota.assert_called_once()
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip_dual")
+def test_run_ota_impl_pauses_after_reaching_device(
+    mock_socket: Mock,
+    firmware_file: Path,
+    mock_perform_ota: Mock,
+    mock_sleep: Mock,
+) -> None:
+    """Test run_ota_impl_ pauses before the next address once the device was reached."""
+    mock_perform_ota.side_effect = [
+        espota2.OTANetworkError("sending data: connection reset"),
+        None,
+    ]
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 0
+    assert result_host == "192.168.1.100"
+    # The first attempt reached the device, so the next one waits first even
+    # though it targets a fresh address
+    mock_sleep.assert_called_once_with(espota2.UPLOAD_RETRY_DELAY)
+
+
+@pytest.mark.usefixtures("mock_socket_constructor", "mock_resolve_ip")
+def test_run_ota_impl_device_error_not_retried(
+    mock_socket: Mock, firmware_file: Path, mock_perform_ota: Mock, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ fails immediately on a device-reported error."""
+    mock_perform_ota.side_effect = espota2.OTAError(
+        "Authentication invalid. Is the password correct?"
+    )
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 1
+    assert result_host is None
+    mock_perform_ota.assert_called_once()
+    mock_sleep.assert_not_called()
+
+
+def test_run_ota_impl_no_addresses(
+    firmware_file: Path, mock_resolve_ip: Mock, mock_sleep: Mock
+) -> None:
+    """Test run_ota_impl_ fails cleanly when resolution yields no addresses."""
+    mock_resolve_ip.return_value = []
+
+    result_code, result_host = espota2.run_ota_impl_(
+        "test.local", 3232, "password", str(firmware_file)
+    )
+
+    assert result_code == 1
+    assert result_host is None
+    mock_sleep.assert_not_called()
 
 
 def test_run_ota_impl_resolve_failed(tmp_path: Path, mock_resolve_ip: Mock) -> None: