From 5535b05bf57c720429e714559be953934bf4c78c Mon Sep 17 00:00:00 2001
From: Mat931 <49403702+Mat931@users.noreply.github.com>
Date: Thu, 30 Apr 2026 19:13:40 +0200
Subject: [PATCH] Fix tests

---
 esphome/espota2.py               |  6 ---
 tests/unit_tests/test_espota2.py | 77 ++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/esphome/espota2.py b/esphome/espota2.py
index ec9cc0f90a..e1cc52e7ce 100644
--- a/esphome/espota2.py
+++ b/esphome/espota2.py
@@ -287,12 +287,6 @@ def perform_ota(
     if ota_type != OTA_TYPE_UPDATE_APP:
         raise OTAError(f"Unsupported OTA type: 0x{ota_type:02X}")
 
-    if (
-        ota_type != OTA_TYPE_UPDATE_APP
-        and not features & SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS
-    ):
-        raise OTAError("Device only supports app updates")
-
     if features & SERVER_FEATURE_SUPPORTS_COMPRESSION:
         upload_contents = gzip.compress(file_contents, compresslevel=9)
         _LOGGER.info("Compressed to %s bytes", len(upload_contents))
diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py
index 7ac88ee129..37b9ef173f 100644
--- a/tests/unit_tests/test_espota2.py
+++ b/tests/unit_tests/test_espota2.py
@@ -776,3 +776,80 @@ def test_perform_ota_version_differences(
 
     # For v2.0, verify more recv calls due to chunk acknowledgments
     assert mock_socket.recv.call_count == 9  # v2.0 has 9 recv calls (includes chunk OK)
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_extended_protocol_app(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test OTA partition table update."""
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_FEATURE_FLAGS]),  # Device supports extended protocol
+        bytes(
+            [
+                espota2.SERVER_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.SERVER_FEATURE_SUPPORTS_PARTITION_ACCESS
+            ]
+        ),  # Device feature flags
+        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
+        bytes([espota2.RESPONSE_CHUNK_OK]),  # Chunk OK
+        bytes([espota2.RESPONSE_RECEIVE_OK]),  # Receive OK
+        bytes([espota2.RESPONSE_UPDATE_END_OK]),  # Update end OK
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    espota2.perform_ota(
+        mock_socket,
+        "testpass",
+        mock_file,
+        "partitions.bin",
+        espota2.OTA_TYPE_UPDATE_APP,
+    )
+
+    # Verify magic bytes were sent
+    assert mock_socket.sendall.call_args_list[0] == call(bytes(espota2.MAGIC_BYTES))
+
+    # Verify features were sent (compression + SHA256 support + extended protocol)
+    assert mock_socket.sendall.call_args_list[1] == call(
+        bytes(
+            [
+                espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
+                | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
+            ]
+        )
+    )
+
+    # Verify ota type was sent
+    assert mock_socket.sendall.call_args_list[2] == call(
+        bytes([espota2.OTA_TYPE_UPDATE_APP])
+    )
+
+
+@pytest.mark.usefixtures("mock_time")
+def test_perform_ota_extended_protocol_unsupported_type(
+    mock_socket: Mock, mock_file: io.BytesIO
+) -> None:
+    """Test OTA fails when OTA type is unsupported by the client."""
+    # Setup socket responses for recv calls
+    recv_responses = [
+        bytes([espota2.RESPONSE_OK]),  # First byte of version response
+        bytes([espota2.OTA_VERSION_2_0]),  # Version number
+        bytes([espota2.RESPONSE_HEADER_OK]),  # Features response
+    ]
+
+    mock_socket.recv.side_effect = recv_responses
+
+    with pytest.raises(espota2.OTAError, match="Unsupported OTA type"):
+        espota2.perform_ota(
+            mock_socket,
+            "testpass",
+            mock_file,
+            "partitions.bin",
+            255,
+        )