diff --git a/selfdrive/athena/athenad.py b/selfdrive/athena/athenad.py index ca3a1bafbb..cc3ab5e95b 100755 --- a/selfdrive/athena/athenad.py +++ b/selfdrive/athena/athenad.py @@ -206,13 +206,17 @@ def retry_upload(tid: int, end_event: threading.Event, increase_count: bool = Tr break -def cb(sm, item, tid, sz: int, cur: int) -> None: +def cb(sm, item, tid, end_event: threading.Event, sz: int, cur: int) -> None: # Abort transfer if connection changed to metered after starting upload + # or if athenad is shutting down to re-connect the websocket sm.update(0) metered = sm['deviceState'].networkMetered if metered and (not item.allow_cellular): raise AbortTransferException + if end_event.is_set(): + raise AbortTransferException + cur_upload_items[tid] = replace(item, progress=cur / sz if sz else 1) @@ -252,7 +256,7 @@ def upload_handler(end_event: threading.Event) -> None: sz = -1 cloudlog.event("athena.upload_handler.upload_start", fn=fn, sz=sz, network_type=network_type, metered=metered, retry_count=item.retry_count) - response = _do_upload(item, partial(cb, sm, item, tid)) + response = _do_upload(item, partial(cb, sm, item, tid, end_event)) if response.status_code not in (200, 201, 401, 403, 412): cloudlog.event("athena.upload_handler.retry", status_code=response.status_code, fn=fn, sz=sz, network_type=network_type, metered=metered) diff --git a/selfdrive/athena/tests/test_athenad.py b/selfdrive/athena/tests/test_athenad.py index 2fecab1b1b..d59058d7e2 100755 --- a/selfdrive/athena/tests/test_athenad.py +++ b/selfdrive/athena/tests/test_athenad.py @@ -233,6 +233,7 @@ class TestAthenadMethods(unittest.TestCase): time.sleep(0.1) # TODO: verify that upload actually succeeded + # TODO: also check that end_event and metered network raises AbortTransferException self.assertEqual(athenad.upload_queue.qsize(), 0) @parameterized.expand([(500, True), (412, False)])