Skip to content

Commit

Permalink
fix: ensure timeout does not raise cancellation (#46)
Browse files Browse the repository at this point in the history
  • Loading branch information
bdraco authored Apr 25, 2024
1 parent c81b1e3 commit 4575fdd
Showing 1 changed file with 14 additions and 7 deletions.
21 changes: 14 additions & 7 deletions src/bluetooth_auto_recovery/recover.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ def __init__(
self.transport: asyncio.Transport | None = None
self.timeout = timeout
self.connection_mode_future = connection_mode_future
self.loop = asyncio.get_running_loop()

def connection_made(self, transport: asyncio.BaseTransport) -> None:
"""Handle connection made."""
Expand All @@ -155,15 +156,21 @@ def data_received(self, data: bytes) -> None:
async def send(self, *args: Any) -> btmgmt_protocol.Response:
"""Send command."""
pkt_objs = btmgmt_protocol.command(*args)
full_pkt = b""
for frame in pkt_objs:
if frame:
full_pkt += frame.octets
self.future = asyncio.Future()
self.future = self.loop.create_future()
assert self.transport is not None # nosec
self.transport.write(full_pkt)
async with asyncio_timeout(self.timeout):
self.transport.write(b"".join(frame.octets for frame in pkt_objs if frame))
cancel_timeout = self.loop.call_later(
self.timeout, self._timeout_future, self.future
)
try:
return await self.future
finally:
cancel_timeout.cancel()
self.future = None

def _timeout_future(self, future: asyncio.Future[btmgmt_protocol.Response]) -> None:
if future and not future.done():
future.set_exception(asyncio.TimeoutError("Timeout waiting for response"))

def connection_lost(self, exc: Exception | None) -> None:
"""Handle connection lost."""
Expand Down

0 comments on commit 4575fdd

Please sign in to comment.