diff --git a/system/athena/athenad.py b/system/athena/athenad.py index 909e99fa71546e..78285f018e3f56 100755 --- a/system/athena/athenad.py +++ b/system/athena/athenad.py @@ -258,13 +258,13 @@ 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, 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) - retry_upload(tid, end_event) - else: - cloudlog.event("athena.upload_handler.success", fn=fn, sz=sz, network_type=network_type, metered=metered) + with _do_upload(item, partial(cb, sm, item, tid, end_event)) as response: + 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) + retry_upload(tid, end_event) + else: + cloudlog.event("athena.upload_handler.success", fn=fn, sz=sz, network_type=network_type, metered=metered) UploadQueueCache.cache(upload_queue) except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, requests.exceptions.SSLError): diff --git a/system/athena/tests/test_athenad.py b/system/athena/tests/test_athenad.py index a4dcbef67a7bd5..a6bfc689306533 100644 --- a/system/athena/tests/test_athenad.py +++ b/system/athena/tests/test_athenad.py @@ -240,7 +240,7 @@ def test_upload_handler(self, host): @with_upload_handler def test_upload_handler_retry(self, mocker, host, status, retry): mock_put = mocker.patch('requests.put') - mock_put.return_value.status_code = status + mock_put.return_value.__enter__.return_value.status_code = status fn = self._create_file('qlog.zst') item = athenad.UploadItem(path=fn, url=f"{host}/qlog.zst", headers={}, created_at=int(time.time()*1000), id='', allow_cellular=True)