diff --git a/tests/unit_tests/test_espota2.py b/tests/unit_tests/test_espota2.py
index b08d4cecf2..ab0ad7034c 100644
--- a/tests/unit_tests/test_espota2.py
+++ b/tests/unit_tests/test_espota2.py
@@ -276,6 +276,7 @@ def test_perform_ota_successful_md5_auth(
             [
                 espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )
@@ -646,6 +647,7 @@ def test_perform_ota_successful_sha256_auth(
             [
                 espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )
@@ -701,6 +703,7 @@ def test_perform_ota_sha256_fallback_to_md5(
             [
                 espota2.CLIENT_FEATURE_SUPPORTS_COMPRESSION
                 | espota2.CLIENT_FEATURE_SUPPORTS_SHA256_AUTH
+                | espota2.CLIENT_FEATURE_SUPPORTS_EXTENDED_PROTOCOL
             ]
         )
     )