From d1f097cf0765788dd4f04563ee6cec299bbd8bfe Mon Sep 17 00:00:00 2001 From: deanlee Date: Wed, 27 Nov 2024 20:39:21 +0800 Subject: [PATCH] use with --- system/athena/athenad.py | 5 +---- system/athena/tests/test_athenad.py | 2 +- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/system/athena/athenad.py b/system/athena/athenad.py index 4b5aa92988fcf2..78285f018e3f56 100755 --- a/system/athena/athenad.py +++ b/system/athena/athenad.py @@ -259,15 +259,12 @@ def upload_handler(end_event: threading.Event) -> None: 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)) - try: + 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) - finally: - response.close() 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)