From 47bb62b875fe2099a7a99482767bbd99b478f9e5 Mon Sep 17 00:00:00 2001
From: Joost Wooning <jwooning@gmail.com>
Date: Tue, 4 Jan 2022 17:01:33 +0100
Subject: [PATCH] athena: methods for adding and cancelling multiple upload
 requests (#23366)

* multiple upload cancel

* multiple uploads athena method

* cleanup

* cleanup

* more cleanup

* isnt used

* fix test

* actually fix test
---
 selfdrive/athena/athenad.py            | 39 +++++++++++++++++---------
 selfdrive/athena/tests/test_athenad.py |  4 +--
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/selfdrive/athena/athenad.py b/selfdrive/athena/athenad.py
index 996ac97f8e23b6..f58013a15da562 100755
--- a/selfdrive/athena/athenad.py
+++ b/selfdrive/athena/athenad.py
@@ -263,20 +263,29 @@ def do_reboot():
 
 @dispatcher.add_method
 def uploadFileToUrl(fn, url, headers):
-  if len(fn) == 0 or fn[0] == '/' or '..' in fn:
-    return 500
-  path = os.path.join(ROOT, fn)
-  if not os.path.exists(path):
-    return 404
+  return uploadFilesToUrls([[fn, url, headers]])
+
 
-  item = UploadItem(path=path, url=url, headers=headers, created_at=int(time.time() * 1000), id=None)
-  upload_id = hashlib.sha1(str(item).encode()).hexdigest()
-  item = item._replace(id=upload_id)
+@dispatcher.add_method
+def uploadFilesToUrls(files_data):
+  items = []
+  for fn, url, headers in files_data:
+    if len(fn) == 0 or fn[0] == '/' or '..' in fn:
+      return 500
+    path = os.path.join(ROOT, fn)
+    if not os.path.exists(path):
+      return 404
+
+    item = UploadItem(path=path, url=url, headers=headers, created_at=int(time.time() * 1000), id=None)
+    upload_id = hashlib.sha1(str(item).encode()).hexdigest()
+    items.append(item._replace(id=upload_id))
+
+  for item in items:
+    upload_queue.put_nowait(item)
 
-  upload_queue.put_nowait(item)
   UploadQueueCache.cache(upload_queue)
 
-  return {"enqueued": 1, "item": item._asdict()}
+  return {"enqueued": len(items), "items": [i._asdict() for i in items]}
 
 
 @dispatcher.add_method
@@ -287,11 +296,15 @@ def listUploadQueue():
 
 @dispatcher.add_method
 def cancelUpload(upload_id):
-  upload_ids = {item.id for item in list(upload_queue.queue)}
-  if upload_id not in upload_ids:
+  if not isinstance(upload_id, list):
+    upload_id = [upload_id]
+
+  uploading_ids = {item.id for item in list(upload_queue.queue)}
+  cancelled_ids = uploading_ids.intersection(upload_id)
+  if len(cancelled_ids) == 0:
     return 404
 
-  cancelled_uploads.add(upload_id)
+  cancelled_uploads.update(cancelled_ids)
   return {"success": 1}
 
 
diff --git a/selfdrive/athena/tests/test_athenad.py b/selfdrive/athena/tests/test_athenad.py
index bf6d3940f9b46d..ae8d2334d39423 100755
--- a/selfdrive/athena/tests/test_athenad.py
+++ b/selfdrive/athena/tests/test_athenad.py
@@ -141,8 +141,8 @@ def test_uploadFileToUrl(self, host):
 
     resp = dispatcher["uploadFileToUrl"]("qlog.bz2", f"{host}/qlog.bz2", {})
     self.assertEqual(resp['enqueued'], 1)
-    self.assertDictContainsSubset({"path": fn, "url": f"{host}/qlog.bz2", "headers": {}}, resp['item'])
-    self.assertIsNotNone(resp['item'].get('id'))
+    self.assertDictContainsSubset({"path": fn, "url": f"{host}/qlog.bz2", "headers": {}}, resp['items'][0])
+    self.assertIsNotNone(resp['items'][0].get('id'))
     self.assertEqual(athenad.upload_queue.qsize(), 1)
 
   @with_http_server