use with

commaai · Nov 27, 2024 · d1f097c · d1f097c
1 parent b34cf42
commit d1f097c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 5 deletions.
diff --git 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
@@ -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)