connections: fix checksum race condition

This fixes a race condition where the nus_handler() could be called
multiple times before returning to send_block() when it is waiting for
the checksum_ready event.

By setting self.expected_checksum = -1 in the nus_handler() callback
itself, we prevent the next call from assuming that the data is also
a checksum.

Issue: #16

Author: dlech

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 ### Changed
 - Print BLE download progress after chunk is complete instead of before.
+### Fixed
+- Fix occasional bad checksum warning when running program via BLE.
 
 ## [1.0.0-alpha.2] - 2021-04-08
 ### Added

pybricksdev/connections.py

@@ -616,16 +616,17 @@ def line_handler(self, line):
             print(line.decode())
 
     def nus_handler(self, sender, data):
-        # If we are currently downloading a program, treat incoming data as checksum.
-        if self.loading:
-            if data[0] == self.expected_checksum:
-                self.checksum_ready.set()
-                self.logger.debug("Correct checksum: {0}".format(data[0]))
-            else:
-                self.logger.warning("Expected checksum {0} but got {1}".format(
-                    self.expected_checksum,
-                    data[0]
-                ))
+        # If we are currently expecting a checksum, validate it and notify the waiter
+        if self.expected_checksum != -1:
+            checksum = data[0]
+            if checksum != self.expected_checksum:
+                raise RuntimeError(
+                    f"Expected checksum {self.expected_checksum} but got {checksum}"
+                )
+
+            self.expected_checksum = -1
+            self.checksum_ready.set()
+            self.logger.debug(f"Correct checksum: {checksum}")
             return
 
         # Store incoming data
@@ -699,9 +700,9 @@ def disconnected_handler(self, _: BleakClient):
             await self.client.start_notify(NUS_TX_UUID, self.nus_handler)
             await self.client.start_notify(PYBRICKS_UUID, self.pybricks_service_handler)
             self.connected = True
-        except BaseException as ex:
+        except:
             self.disconnect()
-            raise ex
+            raise
 
     async def disconnect(self):
         if self.connected:
@@ -722,8 +723,11 @@ async def send_block(self, data):
         try:
             await self.client.write_gatt_char(NUS_RX_UUID, bytearray(data), False)
             await asyncio.wait_for(self.checksum_ready.wait(), timeout=0.5)
-        finally:
+        except:
+            # normally self.expected_checksum = -1 will be called in nus_handler()
+            # but if we timeout or something like that, we need to reset it here
             self.expected_checksum = -1
+            raise
 
     async def run(self, py_path, wait=True, print_output=True):