cli.oad: work around BlueZ disconnect race condition

BlueZ can get in a bad state if the hub disconnects at the same time as
the host. Ideally, this should be fixed in BlueZ but to avoid triggering
the bug, we can wait for the disconnect.

Author: dlech

pybricksdev/ble/oad/control_point.py

@@ -6,6 +6,7 @@
 from typing import AsyncGenerator
 
 from bleak import BleakClient
+from bleak.exc import BleakError
 
 from ._common import OADReturn, SoftwareVersion, oad_uuid
 
@@ -50,7 +51,11 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, *exc_info):
-        await self._client.stop_notify(OAD_CONTROL_POINT_CHAR_UUID)
+        try:
+            await self._client.stop_notify(OAD_CONTROL_POINT_CHAR_UUID)
+        except BleakError:
+            # ignore if already disconnected
+            pass
 
     def _notification_handler(self, sender, data):
         self._queue.put_nowait(data)

pybricksdev/ble/oad/image_identify.py

@@ -11,6 +11,7 @@
 import struct
 
 from bleak import BleakClient
+from bleak.exc import BleakError
 
 from ._common import ImageInfo, OADReturn, SoftwareVersion, oad_uuid
 
@@ -35,7 +36,11 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, *exc_info):
-        await self._client.stop_notify(OAD_IMAGE_IDENTIFY_CHAR_UUID)
+        try:
+            await self._client.stop_notify(OAD_IMAGE_IDENTIFY_CHAR_UUID)
+        except BleakError:
+            # ignore if already disconnected
+            pass
 
     async def validate(
         self,

pybricksdev/cli/oad.py

@@ -62,10 +62,17 @@ async def flash_oad_image(firmware: BinaryIO) -> None:
         print("No OAD device found")
         return
 
+    disconnect_event = asyncio.Event()
+
+    def on_disconnect(_):
+        disconnect_event.set()
+
     # long timeout in case pairing is needed
-    async with asyncio.timeout(60), BleakClient(device) as client, OADImageIdentify(
+    async with asyncio.timeout(60), BleakClient(
+        device, on_disconnect
+    ) as client, OADImageIdentify(client) as image_identify, OADControlPoint(
         client
-    ) as image_identify, OADControlPoint(client) as control_point:
+    ) as control_point:
         image_block = OADImageBlock(client)
 
         print(f"Connected to {device.name}")
@@ -133,6 +140,10 @@ async def flash_oad_image(firmware: BinaryIO) -> None:
         await control_point.enable_oad_image()
         print("Done.")
 
+        # avoid race condition of requesting disconnect while hub is initiating
+        # disconnect itself - this can leave BlueZ in a a bad state
+        await disconnect_event.wait()
+
 
 async def dump_oad_info():
     """