From b31472c569b006dc5bebf8c5471ad385a8267474 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Fri, 24 Nov 2023 08:36:31 -0600
Subject: [PATCH] Refactor noise _error_on_incorrect_preamble to improve
 coverage (#683)

---
 aioesphomeapi/_frame_helper/noise.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/aioesphomeapi/_frame_helper/noise.py b/aioesphomeapi/_frame_helper/noise.py
index 7936f20..8757660 100644
--- a/aioesphomeapi/_frame_helper/noise.py
+++ b/aioesphomeapi/_frame_helper/noise.py
@@ -253,16 +253,15 @@ class APINoiseFrameHelper(APIFrameHelper):
     def _error_on_incorrect_preamble(self, msg: bytes) -> None:
         """Handle an incorrect preamble."""
         explanation = msg[1:].decode()
-        if explanation == "Handshake MAC failure":
-            self._handle_error_and_close(
-                InvalidEncryptionKeyAPIError(
-                    f"{self._log_name}: Invalid encryption key", self._server_name
-                )
+        if explanation != "Handshake MAC failure":
+            exc = HandshakeAPIError(
+                f"{self._log_name}: Handshake failure: {explanation}"
             )
-            return
-        self._handle_error_and_close(
-            HandshakeAPIError(f"{self._log_name}: Handshake failure: {explanation}")
-        )
+        else:
+            exc = InvalidEncryptionKeyAPIError(
+                f"{self._log_name}: Invalid encryption key", self._server_name
+            )
+        self._handle_error_and_close(exc)
 
     def _handle_handshake(self, msg: bytes) -> None:
         if msg[0] != 0: