[core] Add probe_ota_key, a Noise handshake with the device that sends no image (#20601)

esphome/__main__.py

@@ -483,12 +483,9 @@ def has_ota() -> bool:
 
 def has_native_ota() -> bool:
     """Check if native API OTA upload is available (``platform: esphome``)."""
-    if CONF_OTA not in CORE.config:
-        return False
-    return any(
-        ota_item.get(CONF_PLATFORM) == CONF_ESPHOME
-        for ota_item in CORE.config[CONF_OTA]
-    )
+    from esphome.espota2 import esphome_ota_items
+
+    return bool(esphome_ota_items(CORE.config))
 
 
 def has_web_server_ota() -> bool:
@@ -1359,14 +1356,17 @@ def _choose_ota_platform(config: ConfigType, requested: str | None) -> str:
     return CONF_WEB_SERVER
 
 
+def _esphome_ota_conf(config: ConfigType) -> ConfigType | None:
+    """The validated esphome OTA platform block, if configured."""
+    from esphome.espota2 import esphome_ota_items
+
+    return next(iter(esphome_ota_items(config)), None)
+
+
 def _upload_via_native_api(
     config: ConfigType, network_devices: list[str], args: ArgsProtocol
 ) -> tuple[int, str | None]:
-    ota_conf: ConfigType = {}
-    for ota_item in config.get(CONF_OTA, []):
-        if ota_item.get(CONF_PLATFORM) == CONF_ESPHOME:
-            ota_conf = ota_item
-            break
+    ota_conf = _esphome_ota_conf(config) or {}
 
     from esphome import espota2
     from esphome.components.noise import static_encryption_key
@@ -1457,10 +1457,8 @@ def _old_key_notice(config: ConfigType) -> None:
     `upload` may send a stale image that still carries the previous key."""
     from esphome import espota2
 
-    for item in config.get(CONF_OTA, []):
-        if item.get(CONF_PLATFORM) == CONF_ESPHOME and (
-            item.get(CONF_ENCRYPTION) or {}
-        ).get(espota2.CONF_OLD_KEY):
+    for item in espota2.esphome_ota_items(config):
+        if (item.get(CONF_ENCRYPTION) or {}).get(espota2.CONF_OLD_KEY):
             _LOGGER.warning(espota2.OLD_KEY_REMOVE_NOTICE)
 
 
@@ -1468,12 +1466,12 @@ def _upload_via_web_server(
     config: ConfigType, network_devices: list[str], binary: Path
 ) -> tuple[int, str | None]:
     from esphome import web_server_ota
+    from esphome.espota2 import esphome_ota_items
     from esphome.web_server_helpers import get_web_server_connection
 
     if any(
-        ota_item.get(CONF_PLATFORM) == CONF_ESPHOME
-        and ota_item.get(CONF_ENCRYPTION) is not None
-        for ota_item in config.get(CONF_OTA, [])
+        ota_item.get(CONF_ENCRYPTION) is not None
+        for ota_item in esphome_ota_items(config)
     ):
         _LOGGER.warning(
             "This config has OTA encryption, but the web_server OTA path sends "

esphome/espota2.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 import contextlib
+import enum
 import gzip
 import hashlib
 import io
@@ -12,8 +13,10 @@ import socket
 import time
 from typing import Any
 
+from esphome.const import CONF_ESPHOME, CONF_OTA, CONF_PLATFORM
 from esphome.core import EsphomeError
 from esphome.helpers import ProgressBar, resolve_ip_address
+from esphome.types import ConfigType
 
 OTA_TYPE_UPDATE_APP = 0x00
 OTA_TYPE_UPDATE_PARTITION_TABLE = 0x01
@@ -118,6 +121,7 @@ _AUTH_METHODS: dict[int, tuple[Callable[..., Any], int, str]] = {
     RESPONSE_REQUEST_SHA256_AUTH: (hashlib.sha256, 64, "SHA256"),
     RESPONSE_REQUEST_AUTH: (hashlib.md5, 32, "MD5"),
 }
+_AUTH_ANSWERS = (*_AUTH_METHODS, RESPONSE_AUTH_OK)
 
 # Error response code -> human-readable message (without the "Error: " prefix; check_error()
 # prepends it uniformly). Looked up by check_error() to translate a single byte from the device
@@ -224,10 +228,52 @@ class OTAKeyRejected(OTAError):
     """The Noise handshake failed on the key; another key may succeed."""
 
 
+class OTAEncryptionNotOffered(OTAError):
+    """A key is configured but the device cannot encrypt."""
+
+
+class ProbeResult(enum.Enum):
+    """What probe_ota_key found; only ACCEPTED means the device runs the key."""
+
+    ACCEPTED = "accepted"
+    REJECTED = "rejected"  # the device runs another key
+    NOT_OFFERED = "not offered"  # the device cannot encrypt
+    DEVICE_ERROR = "device error"  # answered, but not about the key
+    UNREACHABLE = "unreachable"  # no answer at all
+
+    def __bool__(self) -> bool:
+        return self is ProbeResult.ACCEPTED
+
+
+_PROBE_ANSWERS = (
+    (OTAKeyRejected, ProbeResult.REJECTED),
+    (OTAEncryptionNotOffered, ProbeResult.NOT_OFFERED),
+)
+
+
+def esphome_ota_items(config: ConfigType) -> list[ConfigType]:
+    """The esphome ota entries; a raw ``ota:`` may be a mapping, not a list."""
+    ota = config.get(CONF_OTA) or []
+    items = [ota] if isinstance(ota, dict) else ota
+    return [
+        item
+        for item in items
+        if isinstance(item, dict) and item.get(CONF_PLATFORM) == CONF_ESPHOME
+    ]
+
+
 # Uploader side options under `ota: encryption:`; the ota component imports
 # the names so the upload path never loads the component module
 CONF_ALLOW_PLAINTEXT_UPLOAD = "allow_plaintext_upload"
 CONF_OLD_KEY = "old_key"
+ENCRYPTION_NOT_OFFERED_ADVICE = (
+    "Refusing to send the image in plaintext. The running firmware predates "
+    "ESPHome 2026.9.0 or has no 'api: encryption: key'. With an api key, install "
+    "once without the 'ota: encryption:' block (that build offers encryption), "
+    f"then restore it; otherwise set '{CONF_ALLOW_PLAINTEXT_UPLOAD}: true' under "
+    "'ota: encryption:' for this one install, or flash by serial or the "
+    "web_server OTA platform."
+)
 ALLOW_PLAINTEXT_UPLOAD_NOTICE = (
     f"'{CONF_ALLOW_PLAINTEXT_UPLOAD}' is set; expected once, on the install that "
     "migrates a device which never encrypted. If this device encrypted before, "
@@ -429,6 +475,22 @@ def send_check(
         raise OTANetworkError(f"sending {msg}: {err}") from err
 
 
+def noise_handshake(psk: str, prologue: bytes) -> Any:
+    """A fresh NNpsk0 initiator; a bad key or a missing noise stack raises here."""
+    # Lazy: the noise stack is only imported for an encrypted upload
+    try:
+        from aioesphomeapi.noise import NoiseHandshake
+    except ImportError as err:
+        raise OTAError(
+            "OTA encryption requires a newer aioesphomeapi; update your "
+            "esphome installation (pip install -U esphome) and retry"
+        ) from err
+    try:
+        return NoiseHandshake(psk, prologue)
+    except ValueError as err:
+        raise OTAError(f"Invalid OTA encryption key: {err}") from err
+
+
 class NoiseSocketWrapper:
     """Runs the OTA session inside a Noise (ChaCha20-Poly1305) transport.
 
@@ -439,25 +501,12 @@ class NoiseSocketWrapper:
     """
 
     def __init__(self, sock: socket.socket, psk: str, prologue: bytes) -> None:
-        # Deliberately lazy: the noise stack (noiseprotocol, cryptography) is
-        # only imported when an encrypted upload actually runs.
-        try:
-            from aioesphomeapi.noise import NoiseHandshake
-        except ImportError as err:
-            raise OTAError(
-                "OTA encryption requires a newer aioesphomeapi; update your "
-                "esphome installation (pip install -U esphome) and retry"
-            ) from err
-        # The aioesphomeapi import above already loaded cryptography; bind
-        # the exception once so recv() pays no per-frame import lookup
+        self._handshake = noise_handshake(psk, prologue)
+        # Bound once; recv() must not import per frame
         from cryptography.exceptions import InvalidTag
 
         self._invalid_tag = InvalidTag
         self._sock = sock
-        try:
-            self._handshake = NoiseHandshake(psk, prologue)
-        except ValueError as err:
-            raise OTAError(f"Invalid OTA encryption key: {err}") from err
         self._encrypt = None
         self._decrypt = None
         self._buffer = b""
@@ -564,43 +613,20 @@ class NoiseSocketWrapper:
         return data
 
 
-def perform_ota(
+def _negotiate_session(
     sock: socket.socket,
-    password: str | None,
-    file_handle: io.IOBase,
-    filename: Path,
-    ota_type: int = OTA_TYPE_UPDATE_APP,
-    noise_psk: str | None = None,
+    noise_psk: str | None,
     plaintext_fallback: bool = False,
     allow_plaintext_upload: bool = False,
-) -> None:
-    # Validate up front; an out-of-range value would only surface as a
-    # ValueError deep inside send_check, bypassing OTAError handling
-    if not isinstance(ota_type, int) or not 0 <= ota_type <= 0xFF:
-        raise OTAError(
-            f"Invalid ota_type {ota_type!r}; expected an integer in range 0-255"
-        )
-    if ota_type not in _SUPPORTED_OTA_TYPES:
-        supported = ", ".join(f"0x{t:02X}" for t in sorted(_SUPPORTED_OTA_TYPES))
-        raise OTAError(
-            f"Unsupported OTA type 0x{ota_type:02X}; this ESPHome supports: {supported}"
-        )
-
-    if noise_psk is not None and not noise_psk:
-        raise OTAError(
-            "An empty OTA encryption key was provided; refusing to upload in plaintext"
-        )
-
-    file_contents = file_handle.read()
-    file_size = len(file_contents)
-    _LOGGER.info("Uploading %s (%s bytes)", filename, file_size)
-
+) -> tuple[socket.socket | NoiseSocketWrapper, int, int, bool]:
+    """Negotiate through the Noise handshake; returns the socket to go on with
+    (wrapped once encrypted), version, features, extended."""
     # Enable nodelay, we need it for phase 1
     sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
     send_check(sock, MAGIC_BYTES, "magic bytes")
 
     _, version = receive_exactly(sock, 2, "version", RESPONSE_OK)
-    _LOGGER.info("Connection established; device supports OTA version %s", version)
+    _LOGGER.debug("Device supports OTA version %s", version)
     supported_versions = (OTA_VERSION_1_0, OTA_VERSION_2_0)
     if version not in supported_versions:
         raise OTAError(
@@ -656,15 +682,9 @@ def perform_ota(
             # capture the image (wifi credentials, api key)
             # Remove before 2027.3.0: installing without the block no longer
             # falls back then; advise 'allow_plaintext_upload: true' instead
-            raise OTAError(
+            raise OTAEncryptionNotOffered(
                 "An OTA encryption key is configured but the device did not "
-                "offer encryption; refusing to send the image in plaintext. "
-                "The running firmware predates ESPHome 2026.9.0 or has no "
-                "'api: encryption: key'. With an api key, install once "
-                "without the 'ota: encryption:' block (that build offers "
-                f"encryption), then restore it; otherwise set '{CONF_ALLOW_PLAINTEXT_UPLOAD}: "
-                "true' under 'ota: encryption:' for this one install, or flash by "
-                "serial or the web_server OTA platform."
+                "offer encryption"
             )
     if noise_psk:
         # The prologue binds every negotiation byte both sides saw, so any
@@ -682,9 +702,50 @@ def perform_ota(
         except OTANetworkError as err:
             # A transport fault: retry encrypted before considering plaintext
             raise OTAHandshakeNetworkError(str(err)) from err
-        _LOGGER.info("Encrypted connection established")
+        _LOGGER.debug("Encrypted connection established")
         if allow_plaintext_upload:
             _LOGGER.warning(ALLOW_PLAINTEXT_UPLOAD_REMOVE_WARNING)
+    return sock, version, features, extended_proto
+
+
+def perform_ota(
+    sock: socket.socket,
+    password: str | None,
+    file_handle: io.IOBase,
+    filename: Path,
+    ota_type: int = OTA_TYPE_UPDATE_APP,
+    noise_psk: str | None = None,
+    plaintext_fallback: bool = False,
+    allow_plaintext_upload: bool = False,
+    on_upload_start: Callable[[], None] | None = None,
+) -> None:
+    # Validate up front; an out-of-range value would only surface as a
+    # ValueError deep inside send_check, bypassing OTAError handling
+    if not isinstance(ota_type, int) or not 0 <= ota_type <= 0xFF:
+        raise OTAError(
+            f"Invalid ota_type {ota_type!r}; expected an integer in range 0-255"
+        )
+    if ota_type not in _SUPPORTED_OTA_TYPES:
+        supported = ", ".join(f"0x{t:02X}" for t in sorted(_SUPPORTED_OTA_TYPES))
+        raise OTAError(
+            f"Unsupported OTA type 0x{ota_type:02X}; this ESPHome supports: {supported}"
+        )
+
+    if noise_psk is not None and not noise_psk:
+        raise OTAError(
+            "An empty OTA encryption key was provided; refusing to upload in plaintext"
+        )
+
+    file_contents = file_handle.read()
+    file_size = len(file_contents)
+    _LOGGER.info("Uploading %s (%s bytes)", filename, file_size)
+
+    sock, version, features, extended_proto = _negotiate_session(
+        sock, noise_psk, plaintext_fallback, allow_plaintext_upload
+    )
+    _LOGGER.info("Connection established; device supports OTA version %s", version)
+    if isinstance(sock, NoiseSocketWrapper):
+        _LOGGER.info("Encrypted connection established")
 
     if ota_type != OTA_TYPE_UPDATE_APP:
         # Any non-app OTA type requires the extended protocol and the
@@ -761,18 +822,15 @@ def perform_ota(
         send_check(sock, result, "auth result")
         receive_exactly(sock, 1, "auth result", RESPONSE_AUTH_OK)
 
-    (auth,) = receive_exactly(
-        sock,
-        1,
-        "auth",
-        [RESPONSE_REQUEST_AUTH, RESPONSE_REQUEST_SHA256_AUTH, RESPONSE_AUTH_OK],
-    )
+    (auth,) = receive_exactly(sock, 1, "auth", list(_AUTH_ANSWERS))
 
     if auth != RESPONSE_AUTH_OK:
         hash_func, nonce_size, hash_name = _AUTH_METHODS[auth]
         perform_auth(sock, password, hash_func, nonce_size, hash_name)
 
     _LOGGER.info("Handshake complete")
+    if on_upload_start is not None:
+        on_upload_start()  # the device may commit anything sent from here
 
     sock.settimeout(DATA_PHASE_TIMEOUT)
 
@@ -890,6 +948,146 @@ def perform_ota(
     time.sleep(1)
 
 
+PROBE_TIMEOUT = 120.0
+PROBE_RETRY_DELAY = 3.0
+PROBE_CONNECT_TIMEOUT = 3.0
+PROBE_HANDSHAKE_TIMEOUT = 10.0
+
+
+def _resolve_targets(remote_host: str | list[str], remote_port: int) -> list[Any]:
+    """Every address to try for the host, through the cache; OTAError when
+    there is none."""
+    from esphome.core import CORE
+
+    try:
+        res = resolve_ip_address(
+            remote_host, remote_port, address_cache=CORE.address_cache
+        )
+    except EsphomeError as err:
+        raise OTAError(err) from err
+    if not res:
+        raise OTAError(f"No addresses to connect to for {remote_host}")
+    return res
+
+
+class _DeadlineSocket:
+    """Re-arms the timeout before every operation so the whole negotiation,
+    not each read, ends by the deadline."""
+
+    def __init__(self, sock: socket.socket, deadline: float) -> None:
+        self._sock = sock
+        self._deadline = deadline
+
+    def _arm(self) -> None:
+        remaining = self._deadline - time.monotonic()
+        if remaining <= 0:
+            raise TimeoutError("probe deadline reached")
+        self._sock.settimeout(min(PROBE_HANDSHAKE_TIMEOUT, remaining))
+
+    def recv(self, amount: int) -> bytes:
+        self._arm()
+        return self._sock.recv(amount)
+
+    def sendall(self, data: bytes) -> None:
+        self._arm()
+        self._sock.sendall(data)
+
+    def setsockopt(self, level: int, optname: int, value: int) -> None:
+        self._sock.setsockopt(level, optname, value)
+
+    def close(self) -> None:
+        self._sock.close()
+
+
+def probe_ota_key(
+    remote_host: str | list[str],
+    remote_port: int,
+    noise_psk: str,
+    *,
+    timeout: float = PROBE_TIMEOUT,
+    retry_rejected: bool = True,
+) -> ProbeResult:
+    """Whether the device completes a Noise handshake with the key, no image
+    sent. Retries until the deadline, a rejection too with ``retry_rejected``
+    (the old firmware may answer once more after an upload)."""
+
+    noise_handshake(noise_psk, b"")  # a local fault must not read as the device's no
+    deadline = time.monotonic() + timeout
+
+    def pause(delay: float) -> None:
+        time.sleep(max(0.0, min(delay, deadline - time.monotonic())))
+
+    res: list[Any] = []
+    answered: set[Any] = set()  # addresses whose answer no retry changes
+    last_error = "no time left"
+    outcome = ProbeResult.UNREACHABLE  # a device answer outranks a network error
+    while time.monotonic() < deadline:
+        if not res:
+            # A rebooting device may drop off mDNS for a while; keep asking.
+            # Addresses the caller pre-resolved (the address cache, or a
+            # literal IP) are its to keep; only a hostname is looked up again
+            try:
+                res = _resolve_targets(remote_host, remote_port)
+            except OTAError as err:
+                if outcome is ProbeResult.UNREACHABLE:
+                    last_error = str(err)
+                _LOGGER.debug("Host not resolved yet (%s); retrying", err)
+                pause(PROBE_RETRY_DELAY)
+                continue
+        # One round: every address that has not answered, back to back
+        started = time.monotonic()
+        unreached = 0
+        for af, socktype, _, _, sa in res:
+            if sa in answered:
+                continue
+            if (remaining := deadline - time.monotonic()) <= 0:
+                break
+            sock = socket.socket(af, socktype)
+            # A dead host must not eat the budget
+            sock.settimeout(min(PROBE_CONNECT_TIMEOUT, remaining))
+            with contextlib.closing(sock):
+                try:
+                    sock.connect(sa)
+                    session = _negotiate_session(
+                        _DeadlineSocket(sock, deadline), noise_psk
+                    )[0]
+                    # A password challenge also proves the key: it comes encrypted
+                    receive_exactly(session, 1, "auth", list(_AUTH_ANSWERS))
+                except (OSError, OTANetworkError) as err:
+                    if outcome is ProbeResult.UNREACHABLE:
+                        last_error = f"{sa[0]}: {err}"
+                    unreached += 1
+                except OTAError as err:
+                    # Final in the precheck; a rebooting device may answer
+                    # differently once it is back
+                    found = next(
+                        (r for t, r in _PROBE_ANSWERS if isinstance(err, t)),
+                        ProbeResult.DEVICE_ERROR,
+                    )
+                    # A stale address answering with some other error must
+                    # not hide what the device said about the key
+                    if found is not ProbeResult.DEVICE_ERROR or not any(
+                        outcome is r for _, r in _PROBE_ANSWERS
+                    ):
+                        outcome, last_error = found, f"{sa[0]}: {err}"
+                    if not retry_rejected:
+                        answered.add(sa)
+                else:
+                    _LOGGER.info("Device %s accepted the key", sa[0])
+                    return ProbeResult.ACCEPTED
+            _LOGGER.debug("Key not accepted yet (%s); retrying", last_error)
+        # Every address gets its own final say; a stale one cannot answer
+        # for the device
+        pending = sum(1 for *_, sa in res if sa not in answered)
+        if pending == 0 or time.monotonic() >= deadline:
+            break
+        if unreached == pending:
+            res = []  # no answer this round: the device may have a new address
+        pause(PROBE_RETRY_DELAY - (time.monotonic() - started))
+    _LOGGER.warning("Checking the key failed, %s: %s", outcome.value, last_error)
+    return outcome
+
+
 def run_ota_impl_(
     remote_host: str | list[str],
     remote_port: int,
@@ -900,16 +1098,13 @@ def run_ota_impl_(
     plaintext_fallback: bool = False,
     allow_plaintext_upload: bool = False,
     old_noise_psk: str | None = None,
+    on_upload_start: Callable[[], None] | None = None,
 ) -> tuple[int, str | None]:
     from esphome.core import CORE
 
-    # Handle both single host and list of hosts
     try:
-        # Resolve all hosts at once for parallel DNS resolution
-        res = resolve_ip_address(
-            remote_host, remote_port, address_cache=CORE.address_cache
-        )
-    except EsphomeError as err:
+        res = _resolve_targets(remote_host, remote_port)
+    except OTAError:
         _LOGGER.error(
             "Error resolving IP address of %s. Is it connected to WiFi?",
             remote_host,
@@ -920,11 +1115,7 @@ def run_ota_impl_(
             "(If this error persists, please set a static IP address: "
             "https://esphome.io/components/wifi/#manual-ips)"
         )
-        raise OTAError(err) from err
-
-    if not res:
-        _LOGGER.error("No addresses to connect to for %s", remote_host)
-        return 1, None
+        raise
 
     # Every address is tried at least once and EXTRA_UPLOAD_ATTEMPTS retries
     # are shared across the addresses, cycling through them. Wait before an
@@ -941,6 +1132,13 @@ def run_ota_impl_(
     reconnect = False  # the device closed on purpose; nothing to wait for
     attempt = 0
     encryption = _EncryptionAttempt(noise_psk, plaintext_fallback, old_noise_psk)
+
+    def upload_started() -> None:
+        nonlocal on_upload_start
+        if on_upload_start is not None:
+            on_upload_start()
+            on_upload_start = None  # once per upload, whatever the retries do
+
     while attempt < total_attempts:
         af, socktype, _, _, sa = res[attempt % len(res)]
         if not reconnect and (reached_device or attempt >= len(res)):
@@ -978,6 +1176,7 @@ def run_ota_impl_(
                     encryption.noise_psk,
                     encryption.plaintext_fallback,
                     allow_plaintext_upload=allow_plaintext_upload,
+                    on_upload_start=upload_started,
                 )
             except OTAKeyRejected as err:
                 # The device sent a reject and closed on its own, so reconnect
@@ -1013,6 +1212,8 @@ def run_ota_impl_(
                 # Device-reported error (wrong password, wrong flash size, ...);
                 # retrying cannot succeed, so fail immediately
                 _LOGGER.error(str(err))
+                if isinstance(err, OTAEncryptionNotOffered):
+                    _LOGGER.error(ENCRYPTION_NOT_OFFERED_ADVICE)
                 return 1, None
 
         if encryption.tried_old_key:
@@ -1038,7 +1239,11 @@ def run_ota(
     plaintext_fallback: bool = False,
     allow_plaintext_upload: bool = False,
     old_noise_psk: str | None = None,
+    on_upload_start: Callable[[], None] | None = None,
 ) -> tuple[int, str | None]:
+    """``on_upload_start`` runs once, when the device first accepted the
+    session and the image is about to go out: the point from which it may
+    commit, which a later retry or failure does not undo."""
     try:
         return run_ota_impl_(
             remote_host,
@@ -1050,6 +1255,7 @@ def run_ota(
             plaintext_fallback,
             allow_plaintext_upload,
             old_noise_psk,
+            on_upload_start,
         )
     except OTAError as err:
         _LOGGER.error(err)

tests/unit_tests/test_espota2.py

@@ -964,10 +964,10 @@ def test_run_ota_impl_device_error_not_retried(
 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."""
+    """An empty resolution is reported once and fails cleanly."""
     mock_resolve_ip.return_value = []
 
-    result_code, result_host = espota2.run_ota_impl_(
+    result_code, result_host = espota2.run_ota(
         "test.local", 3232, "password", str(firmware_file)
     )
 

tests/unit_tests/test_espota2_noise.py

@@ -13,6 +13,7 @@ import base64
 from collections.abc import Callable
 import hashlib
 import io
+import itertools
 import logging
 from pathlib import Path
 import socket
@@ -24,11 +25,16 @@ from unittest.mock import Mock, patch
 import pytest
 
 from esphome import espota2
+from esphome.core import EsphomeError
+from esphome.espota2 import ProbeResult
 
 PSK = base64.b64encode(bytes(range(32))).decode()
 OTHER_PSK = base64.b64encode(bytes(range(1, 33))).decode()
 
 MAGIC = bytes(espota2.MAGIC_BYTES)
+# One resolved loopback address, as resolve_ip_address returns it
+RESOLVED = [(2, 1, 0, "", ("127.0.0.1", 1))]
+TWO_ADDRESSES = [(2, 1, 0, "", ("10.0.0.1", 1)), (2, 1, 0, "", ("10.0.0.2", 1))]
 
 
 def _recv_exact(sock: socket.socket, amount: int) -> bytes:
@@ -82,6 +88,7 @@ class FakeEncryptedDevice(threading.Thread):
         self.require_noise = require_noise
         self.prologue_features_override = prologue_features_override
         self.received: bytes | None = None
+        self.probes = 0  # clients that left right after the handshake
         self.error: Exception | None = None
         self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.listener.bind(("127.0.0.1", 0))
@@ -195,7 +202,12 @@ class FakeEncryptedDevice(threading.Thread):
     ) -> None:
         """The post-handshake exchange, identical over both transports."""
         send_byte(espota2.RESPONSE_AUTH_OK)
-        recv_unit(1)  # ota type
+        try:
+            recv_unit(1)  # ota type
+        except ConnectionError:
+            # A key probe leaves here, like the firmware it is patterned on
+            self.probes += 1
+            return
         size = int.from_bytes(recv_unit(4), "big")
         send_byte(espota2.RESPONSE_UPDATE_PREPARE_OK)
         md5_hex = recv_unit(32)
@@ -250,6 +262,7 @@ def _run_ota(
     noise_psk: str,
     plaintext_fallback: bool = True,
     old_noise_psk: str | None = None,
+    **kwargs: Any,
 ) -> int:
     """Drive the retry loop, which is where the fallback and the old key reconnect."""
     path = tmp_path / "firmware.bin"
@@ -263,10 +276,110 @@ def _run_ota(
         noise_psk=noise_psk,
         plaintext_fallback=plaintext_fallback,
         old_noise_psk=old_noise_psk,
+        **kwargs,
     )
     return rc
 
 
+def test_on_upload_start_runs_once_the_device_accepted_the_session(
+    tmp_path: Path,
+) -> None:
+    """The rejected first key never reaches the hook; the old_key reconnect does."""
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice(psk=OTHER_PSK, connections=2)
+    on_upload_start = Mock()
+    with patch("time.sleep"):
+        rc = _run_ota(
+            device,
+            b"firmware",
+            tmp_path,
+            PSK,
+            plaintext_fallback=False,
+            old_noise_psk=OTHER_PSK,
+            on_upload_start=on_upload_start,
+        )
+    device.join_and_check()
+    assert rc == 0
+    on_upload_start.assert_called_once()
+
+
+def test_on_upload_start_fires_once_across_retries() -> None:
+    """A data phase fault after the hook retries the upload; the hook has
+    already reported that the device may commit."""
+    calls: list[int] = []
+
+    def perform(*_args: Any, **kwargs: Any) -> None:
+        kwargs["on_upload_start"]()
+        if len(calls) == 0:
+            calls.append(1)
+            raise espota2.OTANetworkError("reset")
+
+    hook = Mock()
+    with (
+        patch("esphome.espota2.resolve_ip_address", return_value=RESOLVED),
+        patch("socket.socket"),
+        patch("pathlib.Path.open"),
+        patch("esphome.espota2.perform_ota", side_effect=perform),
+        patch("time.sleep"),
+    ):
+        rc, _ = espota2.run_ota_impl_(
+            "h", 1, None, Path("fw.bin"), on_upload_start=hook
+        )
+    assert rc == 0
+    hook.assert_called_once()
+
+
+def test_probe_result_is_true_only_when_accepted() -> None:
+    assert ProbeResult.ACCEPTED
+    assert not any(r for r in ProbeResult if r is not ProbeResult.ACCEPTED)
+
+
+def test_probe_ota_key_keeps_the_key_answer_over_a_stale_device_error(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """A stale address running some other listener answers with its own
+    error; what the device said about the key still wins."""
+    with (
+        patch("esphome.espota2.resolve_ip_address", return_value=TWO_ADDRESSES),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[
+                espota2.OTAKeyRejected("no"),
+                espota2.OTAError("unsupported OTA version 9"),
+            ],
+        ),
+        patch("time.sleep"),
+        caplog.at_level(logging.WARNING),
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.REJECTED
+        )
+    assert "failed, rejected: 10.0.0.1: no" in caplog.text
+
+
+def test_on_upload_start_failure_closes_the_socket() -> None:
+    def start_upload(*_args: Any, **kwargs: Any) -> None:
+        kwargs["on_upload_start"]()
+
+    with (
+        patch("esphome.espota2.resolve_ip_address", return_value=RESOLVED),
+        patch("socket.socket") as sock,
+        patch("pathlib.Path.open"),
+        patch("esphome.espota2.perform_ota", side_effect=start_upload),
+        pytest.raises(RuntimeError, match="boom"),
+    ):
+        espota2.run_ota_impl_(
+            "h",
+            1,
+            None,
+            Path("fw.bin"),
+            on_upload_start=Mock(side_effect=RuntimeError("boom")),
+        )
+    sock.return_value.close.assert_called_once()
+
+
 THIRD_PSK = base64.b64encode(bytes(range(2, 34))).decode()
 
 
@@ -363,9 +476,7 @@ def test_tampered_negotiation_breaks_handshake() -> None:
 def test_client_fails_closed_when_device_lacks_encryption() -> None:
     """With a key configured, a device not offering noise aborts the upload."""
     device = FakeEncryptedDevice(offer_noise=False, require_noise=False)
-    with pytest.raises(
-        espota2.OTAError, match="refusing to send the image.*allow_plaintext_upload"
-    ):
+    with pytest.raises(espota2.OTAEncryptionNotOffered, match="did not offer"):
         _upload(device, b"firmware", PSK)
     device.join_and_check()
 
@@ -666,6 +777,120 @@ def test_non_key_reject_reason_is_a_device_error(
     assert not any("Retrying in plaintext" in r.message for r in caplog.records)
 
 
+def test_probe_ota_key_accepts_the_running_key(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """The negotiation's upload-style progress lines stay below INFO."""
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice()
+    device.start()
+    with caplog.at_level(logging.INFO):
+        assert (
+            espota2.probe_ota_key("127.0.0.1", device.port, PSK, timeout=5)
+            is ProbeResult.ACCEPTED
+        )
+    device.join_and_check()
+    assert device.probes == 1
+    assert device.received is None
+    assert "accepted the key" in caplog.text
+    assert "Connection established" not in caplog.text
+
+
+@pytest.mark.parametrize(
+    "device_kwargs",
+    [{"psk": OTHER_PSK}, {"offer_noise": False, "require_noise": False}],
+    ids=["wrong_key", "no_offer"],
+)
+def test_probe_ota_key_gives_up_at_the_deadline(
+    device_kwargs: dict[str, Any], caplog: pytest.LogCaptureFixture
+) -> None:
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice(connections=2, **device_kwargs)
+    device.start()
+    with patch("time.sleep"), caplog.at_level(logging.WARNING):
+        assert (
+            espota2.probe_ota_key("127.0.0.1", device.port, PSK, timeout=0.5)
+            is not ProbeResult.ACCEPTED
+        )
+    device.join_and_check()
+    assert any("Checking the key failed" in r.message for r in caplog.records)
+    assert device.received is None
+
+
+def test_probe_ota_key_fails_fast_on_a_definitive_answer() -> None:
+    """A device in steady state that rejects the key gives its final answer
+    on the first attempt; the precheck does not wait out the deadline."""
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice(psk=OTHER_PSK)
+    device.start()
+    with patch("time.sleep") as sleep:
+        assert (
+            espota2.probe_ota_key(
+                "127.0.0.1", device.port, PSK, timeout=30, retry_rejected=False
+            )
+            is ProbeResult.REJECTED
+        )
+    device.join_and_check()
+    sleep.assert_not_called()
+
+
+def test_probe_ota_key_retries_after_a_transport_fault() -> None:
+    """The device may still be rebooting right after the upload."""
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice(connections=2, drop_handshakes=1)
+    device.start()
+    with patch("time.sleep"):
+        assert (
+            espota2.probe_ota_key("127.0.0.1", device.port, PSK, timeout=5)
+            is ProbeResult.ACCEPTED
+        )
+    device.join_and_check()
+    assert device.probes == 1
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{"side_effect": EsphomeError("no such host")}, {"return_value": []}],
+    ids=["error", "empty"],
+)
+def test_probe_ota_key_resolution_failure(
+    kwargs: dict[str, Any], caplog: pytest.LogCaptureFixture
+) -> None:
+    """A name that never resolves is retried until the deadline, then reported."""
+    pytest.importorskip("aioesphomeapi.noise")
+    with (
+        patch("esphome.espota2.resolve_ip_address", **kwargs) as resolve,
+        patch("time.sleep"),
+        caplog.at_level(logging.WARNING),
+    ):
+        assert (
+            espota2.probe_ota_key("nowhere.local", 3232, PSK, timeout=0.2)
+            is not ProbeResult.ACCEPTED
+        )
+    assert resolve.call_count > 1
+    assert any("failed, unreachable" in r.message for r in caplog.records)
+
+
+def test_probe_ota_key_resolves_again_after_a_reboot() -> None:
+    """The lookup fails while the device reboots, then the probe goes on."""
+    pytest.importorskip("aioesphomeapi.noise")
+    device = FakeEncryptedDevice()
+    device.start()
+    real = espota2.resolve_ip_address
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            side_effect=[EsphomeError("not yet"), real("127.0.0.1", device.port)],
+        ),
+        patch("time.sleep"),
+    ):
+        assert (
+            espota2.probe_ota_key("127.0.0.1", device.port, PSK, timeout=5)
+            is ProbeResult.ACCEPTED
+        )
+    device.join_and_check()
+
+
 def test_bare_block_refuses_a_device_that_cannot_encrypt(
     caplog: pytest.LogCaptureFixture, tmp_path: Path
 ) -> None:
@@ -677,5 +902,399 @@ def test_bare_block_refuses_a_device_that_cannot_encrypt(
     device.join_and_check()
     assert rc == 1
     assert device.received != b"firmware"
-    assert any("refusing to send the image" in r.message for r in caplog.records)
+    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)
+
+
+def test_probe_ota_key_recomputes_the_budget_after_connect() -> None:
+    """A slow connect leaves less for the handshake than the budget had."""
+    clock = [0.0]
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=RESOLVED,
+        ),
+        patch("socket.socket") as sock_cls,
+        patch("time.sleep"),
+        patch("time.monotonic", side_effect=lambda: clock[0]),
+    ):
+        sock = sock_cls.return_value
+        sock.connect.side_effect = lambda _sa: clock.__setitem__(0, clock[0] + 0.4)
+        sock.recv.return_value = b""
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=0.5) is not ProbeResult.ACCEPTED
+        )
+    connect_timeout, handshake_timeout = (
+        c.args[0] for c in sock.settimeout.call_args_list[:2]
+    )
+    assert connect_timeout == 0.5
+    assert handshake_timeout == pytest.approx(0.1)
+
+
+def test_probe_ota_key_tries_every_address_before_a_final_no() -> None:
+    """With retry_rejected off a rejection is final only once every resolved
+    address has answered; a stale cached IP must not veto the real device."""
+    session = Mock()
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=TWO_ADDRESSES,
+        ),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[espota2.OTAKeyRejected("no"), (session, 2, 0, True)],
+        ) as negotiate,
+        patch("esphome.espota2.receive_exactly"),
+        patch("time.sleep") as sleep,
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.ACCEPTED
+        )
+    assert negotiate.call_count == 2
+    sleep.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "answer", [espota2.RESPONSE_REQUEST_AUTH, espota2.RESPONSE_REQUEST_SHA256_AUTH]
+)
+def test_probe_ota_key_takes_a_password_challenge_as_proof(answer: int) -> None:
+    """A device with a password as well as encryption answers the handshake
+    with a challenge; it arrives encrypted, so the key was accepted."""
+    session = Mock()
+    session.recv.return_value = bytes([answer])
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=RESOLVED,
+        ),
+        patch("socket.socket"),
+        patch("esphome.espota2._negotiate_session", return_value=(session, 2, 0, True)),
+    ):
+        assert espota2.probe_ota_key("h", 1, PSK, timeout=30) is ProbeResult.ACCEPTED
+
+
+def test_probe_ota_key_waits_for_every_address_to_answer() -> None:
+    """One address rejecting does not stand in for another that only failed
+    to connect so far; the one that answered is not asked again."""
+    session = Mock()
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=TWO_ADDRESSES,
+        ),
+        patch("socket.socket") as sock,
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[
+                espota2.OTAKeyRejected("no"),
+                OSError("refused"),
+                (session, 2, 0, True),
+            ],
+        ) as negotiate,
+        patch("esphome.espota2.receive_exactly"),
+        patch("time.sleep"),
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.ACCEPTED
+        )
+    assert negotiate.call_count == 3
+    assert [c.args[0] for c in sock.return_value.connect.call_args_list] == [
+        ("10.0.0.1", 1),
+        ("10.0.0.2", 1),
+        ("10.0.0.2", 1),
+    ]
+
+
+def test_probe_ota_key_pauses_between_rounds_once_an_address_answered() -> None:
+    """A stale address before a device that said no keeps the round pacing:
+    the device is not asked again and the stale address is not hammered."""
+    clock = [0.0]
+
+    def sleep(seconds: float) -> None:
+        clock[0] += seconds
+
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=TWO_ADDRESSES,
+        ) as resolve,
+        patch("socket.socket") as sock_cls,
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=itertools.chain(
+                [OSError("refused"), espota2.OTAKeyRejected("no")],
+                itertools.repeat(OSError("refused")),
+            ),
+        ) as negotiate,
+        patch("time.sleep", side_effect=sleep) as sleeper,
+        patch("time.monotonic", side_effect=lambda: clock[0]),
+    ):
+        assert ProbeResult.ACCEPTED is not espota2.probe_ota_key(
+            "h", 1, PSK, timeout=30, retry_rejected=False
+        )
+    # Ten rounds of one stale attempt each, a pause between every two
+    assert negotiate.call_count == 11
+    assert sleeper.call_count == 10
+    assert [c.args[0] for c in sock_cls.return_value.connect.call_args_list].count(
+        ("10.0.0.2", 1)
+    ) == 1
+    # Every later round had no answer, so the host was resolved again each time
+    assert resolve.call_count == 10
+
+
+def test_probe_ota_key_keeps_trying_an_address_a_new_lookup_added() -> None:
+    """An address that answered earlier and is gone from the lookup does not
+    count against the ones still pending."""
+    clock = [0.0]
+
+    def sleep(seconds: float) -> None:
+        clock[0] += seconds
+
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            side_effect=itertools.chain(
+                [TWO_ADDRESSES],
+                itertools.repeat([(2, 1, 0, "", ("10.0.0.3", 1))]),
+            ),
+        ),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=itertools.chain(
+                [espota2.OTAKeyRejected("no")], itertools.repeat(OSError("down"))
+            ),
+        ) as negotiate,
+        patch("time.sleep", side_effect=sleep),
+        patch("time.monotonic", side_effect=lambda: clock[0]),
+    ):
+        assert ProbeResult.ACCEPTED is not espota2.probe_ota_key(
+            "h", 1, PSK, timeout=30, retry_rejected=False
+        )
+    assert negotiate.call_count == 11  # the new address was tried every round
+
+
+def test_probe_ota_key_keeps_the_device_answer_when_a_later_lookup_fails(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    clock = [0.0]
+
+    def sleep(seconds: float) -> None:
+        clock[0] += seconds
+
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            side_effect=itertools.chain(
+                [TWO_ADDRESSES], itertools.repeat(EsphomeError("dns down"))
+            ),
+        ),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[espota2.OTAKeyRejected("no"), OSError("down")],
+        ),
+        patch("time.sleep", side_effect=sleep),
+        patch("time.monotonic", side_effect=lambda: clock[0]),
+        caplog.at_level(logging.WARNING),
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.REJECTED
+        )
+    assert "failed, rejected: 10.0.0.1: no" in caplog.text
+
+
+def test_probe_ota_key_names_a_missing_encryption_offer(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """The upload's advice about the image does not apply to a probe."""
+    with (
+        patch("esphome.espota2.resolve_ip_address", return_value=RESOLVED),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=espota2.OTAEncryptionNotOffered("did not offer encryption"),
+        ),
+        patch("time.sleep"),
+        caplog.at_level(logging.WARNING),
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.NOT_OFFERED
+        )
+    assert "failed, not offered: 127.0.0.1: did not offer encryption" in caplog.text
+    assert "image" not in caplog.text
+
+
+def test_probe_ota_key_resolves_again_after_a_round_without_an_answer() -> None:
+    """A round with no answer looks the host up again, and the device is
+    found under its new address."""
+    session = Mock()
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            side_effect=[
+                [(2, 1, 0, "", ("10.0.0.1", 1))],
+                [(2, 1, 0, "", ("10.0.0.2", 1))],
+            ],
+        ) as resolve,
+        patch("socket.socket") as sock,
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[OSError("unreachable"), (session, 2, 0, True)],
+        ),
+        patch("esphome.espota2.receive_exactly"),
+        patch("time.sleep"),
+    ):
+        assert espota2.probe_ota_key("h", 1, PSK, timeout=30) is ProbeResult.ACCEPTED
+    assert resolve.call_count == 2
+    assert [c.args[0] for c in sock.return_value.connect.call_args_list] == [
+        ("10.0.0.1", 1),
+        ("10.0.0.2", 1),
+    ]
+
+
+def test_probe_ota_key_reports_the_device_answer_over_a_network_error(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """A stale address failing last must not hide that the device said no."""
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=TWO_ADDRESSES,
+        ),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=itertools.chain(
+                [espota2.OTAKeyRejected("wrong key")],
+                itertools.repeat(OSError("refused")),
+            ),
+        ),
+        patch("time.sleep"),
+        caplog.at_level(logging.WARNING),
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=0.05, retry_rejected=False)
+            is ProbeResult.REJECTED
+        )
+    assert "failed, rejected: 10.0.0.1: wrong key" in caplog.text
+
+
+def test_probe_ota_key_reports_a_local_fault_as_its_own() -> None:
+    """A malformed key never reads as the device refusing it."""
+    pytest.importorskip("aioesphomeapi.noise")
+    with (
+        patch("esphome.espota2.resolve_ip_address") as resolve,
+        pytest.raises(espota2.OTAError, match="Invalid OTA encryption key"),
+    ):
+        espota2.probe_ota_key("h", 1, "not-base64!!!", timeout=30)
+    resolve.assert_not_called()
+
+
+def test_probe_ota_key_takes_another_device_answer_as_final() -> None:
+    """An unsupported protocol version does not change with a retry."""
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=RESOLVED,
+        ),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=espota2.OTAError("unsupported OTA version 9"),
+        ) as negotiate,
+        patch("time.sleep") as sleep,
+    ):
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=30, retry_rejected=False)
+            is ProbeResult.DEVICE_ERROR
+        )
+    negotiate.assert_called_once()
+    sleep.assert_not_called()
+
+
+def test_probe_ota_key_retries_a_device_error_while_it_reboots() -> None:
+    """After an upload the device may answer oddly until it is back."""
+    session = Mock()
+    with (
+        patch("esphome.espota2.resolve_ip_address", return_value=RESOLVED),
+        patch("socket.socket"),
+        patch(
+            "esphome.espota2._negotiate_session",
+            side_effect=[espota2.OTAError("busy"), (session, 2, 0, True)],
+        ) as negotiate,
+        patch("esphome.espota2.receive_exactly"),
+        patch("time.sleep"),
+    ):
+        assert espota2.probe_ota_key("h", 1, PSK, timeout=30) is ProbeResult.ACCEPTED
+    assert negotiate.call_count == 2
+
+
+def test_probe_ota_key_deadline_spans_the_whole_negotiation() -> None:
+    """Each read re-arms the timeout from what is left, so a slow peer
+    cannot stretch the probe one full timeout per read."""
+    clock = [0.0]
+    sock = Mock()
+
+    def slow_recv(_amount: int) -> bytes:
+        clock[0] += 0.3
+        return b"\x00\x02"
+
+    sock.recv.side_effect = slow_recv
+    with patch("time.monotonic", side_effect=lambda: clock[0]):
+        wrapped = espota2._DeadlineSocket(sock, 0.5)
+        wrapped.sendall(b"x")
+        assert wrapped.recv(2) == b"\x00\x02"
+        assert wrapped.recv(2) == b"\x00\x02"
+        with pytest.raises(TimeoutError):
+            wrapped.recv(1)
+    timeouts = [c.args[0] for c in sock.settimeout.call_args_list]
+    assert timeouts == [0.5, 0.5, pytest.approx(0.2)]
+
+
+def test_probe_ota_key_stops_when_connect_used_the_budget() -> None:
+    """No handshake timeout is granted past the deadline."""
+    clock = [0.0]
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=RESOLVED,
+        ),
+        patch("socket.socket") as sock_cls,
+        patch("time.sleep") as sleep,
+        patch("time.monotonic", side_effect=lambda: clock[0]),
+    ):
+        sock = sock_cls.return_value
+        sock.connect.side_effect = lambda _sa: clock.__setitem__(0, clock[0] + 1.0)
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=0.5) is not ProbeResult.ACCEPTED
+        )
+    assert [c.args[0] for c in sock.settimeout.call_args_list] == [0.5]
+    sock.sendall.assert_not_called()
+    sleep.assert_not_called()
+
+
+def test_probe_ota_key_never_waits_past_its_deadline() -> None:
+    """Connect and handshake timeouts and the retry sleep are all capped by
+    what is left of the budget."""
+    with (
+        patch(
+            "esphome.espota2.resolve_ip_address",
+            return_value=RESOLVED,
+        ),
+        patch("socket.socket") as sock_cls,
+        patch("time.sleep") as sleep,
+    ):
+        sock = sock_cls.return_value
+        sock.connect.side_effect = OSError("refused")
+        assert (
+            espota2.probe_ota_key("h", 1, PSK, timeout=0.5) is not ProbeResult.ACCEPTED
+        )
+    timeouts = [c.args[0] for c in sock.settimeout.call_args_list]
+    assert all(t <= 0.5 for t in timeouts)
+    assert all(c.args[0] <= 0.5 for c in sleep.call_args_list)