diff --git a/src/infuse_iot/tools/ota_upgrade.py b/src/infuse_iot/tools/ota_upgrade.py index ed1a661..54a8a8c 100644 --- a/src/infuse_iot/tools/ota_upgrade.py +++ b/src/infuse_iot/tools/ota_upgrade.py @@ -6,6 +6,7 @@ __copyright__ = "Copyright 2024, Embeint Inc" import binascii +import time from rich.live import Live from rich.progress import ( @@ -37,13 +38,16 @@ def __init__(self, args): self._client = LocalClient(default_multicast_address(), 60.0) self._min_rssi: int | None = args.rssi self._release: ValidRelease = args.release + self._app_name = self._release.metadata["application"]["primary"] self._app_id = self._release.metadata["application"]["id"] self._new_ver = self._release.metadata["application"]["version"] self._handled: list[int] = [] + self._pending: dict[int, float] = {} self._missing_diffs: set[str] = set() self._already = 0 self._updated = 0 self._no_diff = 0 + self._failed = 0 self.patch_file = b"" self.state = "Scanning" self.progress = Progress( @@ -62,10 +66,12 @@ def add_parser(cls, parser): def progress_table(self): table = Table() - table.add_column(self._new_ver) + table.add_column(f"{self._app_name}\n{self._new_ver}") table.add_column("Count") table.add_row("Updated", str(self._updated)) + table.add_row("Pending", str(len(self._pending))) table.add_row("Already", str(self._already)) + table.add_row("Failed", str(self._failed)) table.add_row("No Diff", str(self._no_diff)) if len(self._missing_diffs) > 0: @@ -93,6 +99,7 @@ def data_progress_cb(self, offset): def run(self): with Live(self.progress_table(), refresh_per_second=4) as live: for source, announce in self._client.observe_announce(): + self.state_update(live, "Scanning") if announce.application != self._app_id: continue if source.infuse_id in self._handled: @@ -100,6 +107,19 @@ def run(self): v = announce.version v_str = f"{v.major}.{v.minor}.{v.revision}+{v.build_num:08x}" + # Check against pending upgrades + if source.infuse_id in self._pending: + if time.time() < self._pending[source.infuse_id]: + # Device could still be applying the upgrade + continue + self._pending.pop(source.infuse_id) + self._handled.append(source.infuse_id) + if v_str == self._new_ver: + self._updated += 1 + else: + self._failed += 1 + continue + # Already running the requested version? if v_str == self._new_ver: self._handled.append(source.infuse_id) @@ -147,8 +167,7 @@ def run(self): ) if hdr.return_code == 0: - self._handled.append(source.infuse_id) - self._updated += 1 + self._pending[source.infuse_id] = time.time() + 30 except ConnectionRefusedError: self.state_update(live, "Scanning")