diff --git a/aioesphomeapi/client.py b/aioesphomeapi/client.py index 057e616..df08de0 100644 --- a/aioesphomeapi/client.py +++ b/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: diff --git a/aioesphomeapi/connection.pxd b/aioesphomeapi/connection.pxd index 31c7997..dfb8694 100644 --- a/aioesphomeapi/connection.pxd +++ b/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) diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py index fac3ea2..4245c7c 100644 --- a/aioesphomeapi/connection.py +++ b/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", diff --git a/tests/test_client.py b/tests/test_client.py index 82e4c7b..d50b001 100644 --- a/tests/test_client.py +++ b/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): diff --git a/tests/test_connection.py b/tests/test_connection.py index 5531624..d1f4b6e 100644 --- a/tests/test_connection.py +++ b/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()