Small cleanups to connection (#648)

aioesphomeapi/client.py

@@ -451,7 +451,7 @@ class APIClient:
 
         assert self._connection is not None
         resp = await self._connection.send_messages_await_response_complex(
-            (ListEntitiesRequest(),), do_append, do_stop, msg_types, timeout=60
+            (ListEntitiesRequest(),), do_append, do_stop, msg_types, 60
         )
         entities: list[EntityInfo] = []
         services: list[UserService] = []
@@ -567,7 +567,7 @@ class APIClient:
 
         message_filter = partial(self._filter_bluetooth_message, address, handle)
         resp = await self._connection.send_messages_await_response_complex(
-            (request,), message_filter, message_filter, msg_types, timeout=timeout
+            (request,), message_filter, message_filter, msg_types, timeout
         )
 
         if isinstance(resp[0], BluetoothGATTErrorResponse):
@@ -893,7 +893,7 @@ class APIClient:
             do_append,
             do_stop,
             msg_types,
-            timeout=DEFAULT_BLE_TIMEOUT,
+            DEFAULT_BLE_TIMEOUT,
         )
         services = []
         for msg in resp:

aioesphomeapi/connection.pxd

@@ -33,7 +33,7 @@ cdef object partial
 cdef object hr
 
 cdef object RESOLVE_TIMEOUT
-cdef object CONNECT_AND_SETUP_TIMEOUT
+cdef object CONNECT_AND_SETUP_TIMEOUT, CONNECT_REQUEST_TIMEOUT
 
 cdef object APIConnectionError
 cdef object BadNameAPIError
@@ -42,10 +42,23 @@ cdef object PingFailedAPIError
 cdef object ReadFailedAPIError
 cdef object TimeoutAPIError
 
+cdef object in_do_connect, astuple
+
+
+@cython.dataclasses.dataclass
+cdef class ConnectionParams:
+    cdef public str address
+    cdef public object port
+    cdef public object password
+    cdef public object client_info
+    cdef public object keepalive
+    cdef public object zeroconf_manager
+    cdef public object noise_psk
+    cdef public object expected_name
 
 cdef class APIConnection:
 
-    cdef object _params
+    cdef ConnectionParams _params
     cdef public object on_stop
     cdef object _on_stop_task
     cdef public object _socket
@@ -95,3 +108,9 @@ cdef class APIConnection:
 
     @cython.locals(handlers=set)
     cpdef _remove_message_callback(self, object on_message, tuple msg_types)
+
+    cpdef _handle_disconnect_request_internal(self, object msg)
+
+    cpdef _handle_ping_request_internal(self, object msg)
+
+    cpdef _handle_get_time_request_internal(self, object msg)

aioesphomeapi/connection.py

@@ -339,7 +339,8 @@ class APIConnection:
         """Step 3 in connect process: initialize the frame helper and init read loop."""
         fh: APIPlaintextFrameHelper | APINoiseFrameHelper
         loop = self._loop
-        assert self._socket is not None
+        if TYPE_CHECKING:
+            assert self._socket is not None
 
         if (noise_psk := self._params.noise_psk) is None:
             _, fh = await loop.create_connection(  # type: ignore[type-var]
@@ -461,7 +462,7 @@ class APIConnection:
         now = loop.time()
 
         if self._send_pending_ping:
-            self.send_message(PING_REQUEST_MESSAGE)
+            self.send_messages((PING_REQUEST_MESSAGE,))
             if self._pong_timer is None:
                 # Do not reset the timer if it's already set
                 # since the only thing we want to reset the timer
@@ -694,7 +695,7 @@ class APIConnection:
         msg_types: tuple[type[Any], ...],
     ) -> Callable[[], None]:
         """Send a message to the remote and register the given message handler."""
-        self.send_message(send_msg)
+        self.send_messages((send_msg,))
         # Since we do not return control to the event loop (no awaits)
         # between sending the message and registering the handler
         # we can be sure that we will not miss any messages even though
@@ -729,7 +730,7 @@ class APIConnection:
         do_append: Callable[[message.Message], bool] | None,
         do_stop: Callable[[message.Message], bool] | None,
         msg_types: tuple[type[Any], ...],
-        timeout: float = 10.0,
+        timeout: _float,
     ) -> list[message.Message]:
         """Send a message to the remote and build up a list response.
 
@@ -779,7 +780,7 @@ class APIConnection:
         return responses
 
     async def send_message_await_response(
-        self, send_msg: message.Message, response_type: Any, timeout: float = 10.0
+        self, send_msg: message.Message, response_type: Any, timeout: _float = 10.0
     ) -> Any:
         [response] = await self.send_messages_await_response_complex(
             (send_msg,),
@@ -887,14 +888,14 @@ class APIConnection:
         # the response if for some reason sending the response
         # fails we will still mark the disconnect as expected
         self._expected_disconnect = True
-        self.send_message(DISCONNECT_RESPONSE_MESSAGE)
+        self.send_messages((DISCONNECT_RESPONSE_MESSAGE,))
         self._cleanup()
 
     def _handle_ping_request_internal(  # pylint: disable=unused-argument
         self, _msg: PingRequest
     ) -> None:
         """Handle a PingRequest."""
-        self.send_message(PING_RESPONSE_MESSAGE)
+        self.send_messages((PING_RESPONSE_MESSAGE,))
 
     def _handle_get_time_request_internal(  # pylint: disable=unused-argument
         self, _msg: GetTimeRequest
@@ -902,7 +903,7 @@ class APIConnection:
         """Handle a GetTimeRequest."""
         resp = GetTimeResponse()
         resp.epoch_seconds = int(time.time())
-        self.send_message(resp)
+        self.send_messages((resp,))
 
     async def disconnect(self) -> None:
         """Disconnect from the API."""
@@ -946,7 +947,7 @@ class APIConnection:
             # Still try to tell the esp to disconnect gracefully
             # but don't wait for it to finish
             try:
-                self.send_message(DISCONNECT_REQUEST_MESSAGE)
+                self.send_messages((DISCONNECT_REQUEST_MESSAGE,))
             except APIConnectionError as err:
                 _LOGGER.error(
                     "%s: Failed to send (forced) disconnect request: %s",

tests/test_client.py

@@ -79,7 +79,7 @@ def auth_client():
 
 
 def patch_response_complex(client: APIClient, messages):
-    async def patched(req, app, stop, msg_types, timeout=5.0):
+    async def patched(req, app, stop, msg_types, timeout):
         resp = []
         for msg in messages:
             if app(msg):

tests/test_connection.py

@@ -97,7 +97,7 @@ async def test_timeout_sending_message(
 
     with pytest.raises(TimeoutAPIError):
         await conn.send_messages_await_response_complex(
-            (PingRequest(),), None, None, (PingResponse,), timeout=0
+            (PingRequest(),), None, None, (PingResponse,), 0
         )
 
     transport.reset_mock()