Small readability improvements to the plain_text frame helper (#626)

esphome · Nov 9, 2023 · 419bf67 · 419bf67
1 parent 94157a8
commit 419bf67
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 12 deletions.
diff --git a/aioesphomeapi/_frame_helper/plain_text.pxd b/aioesphomeapi/_frame_helper/plain_text.pxd
@@ -26,3 +26,5 @@ cdef class APIPlaintextFrameHelper(APIFrameHelper):
         maybe_msg_type=cython.uint
     )
     cpdef data_received(self, bytes data)
+
+    cpdef _error_on_incorrect_preamble(self, object preamble)
diff --git a/aioesphomeapi/_frame_helper/plain_text.py b/aioesphomeapi/_frame_helper/plain_text.py
@@ -95,18 +95,7 @@ def data_received(  # pylint: disable=too-many-branches,too-many-return-statemen
             length_high = init_bytes[1]
             maybe_msg_type = init_bytes[2]
             if preamble != 0x00:
-                if preamble == 0x01:
-                    self._handle_error_and_close(
-                        RequiresEncryptionAPIError(
-                            f"{self._log_name}: Connection requires encryption"
-                        )
-                    )
-                    return
-                self._handle_error_and_close(
-                    ProtocolAPIError(
-                        f"{self._log_name}: Invalid preamble {preamble:02x}"
-                    )
-                )
+                self._error_on_incorrect_preamble(preamble)
                 return
 
             if length_high & 0x80 != 0x80:
@@ -167,3 +156,16 @@ def data_received(  # pylint: disable=too-many-branches,too-many-return-statemen
             self._remove_from_buffer()
             self._on_pkt(msg_type_int, packet_data)
             # If we have more data, continue processing
+
+    def _error_on_incorrect_preamble(self, preamble: _int) -> None:
+        """Handle an incorrect preamble."""
+        if preamble == 0x01:
+            self._handle_error_and_close(
+                RequiresEncryptionAPIError(
+                    f"{self._log_name}: Connection requires encryption"
+                )
+            )
+            return
+        self._handle_error_and_close(
+            ProtocolAPIError(f"{self._log_name}: Invalid preamble {preamble:02x}")
+        )