Attachment db improvements

Framework/MainDriverApi.py

@@ -1396,28 +1396,42 @@ def download_attachments(testcase_info):
 
     # Test case attachments
     for attachment in testcase_info["attachments"]:
-        got_path = db.exists(attachment["hash"])
-        if got_path is None:
-            urls.append({
-                "url": url_prefix + attachment["path"],
-                "download_dir": attachment_path,
-                "attachment": attachment,
-            })
+        entry = db.exists(attachment["hash"])
+        to_append = {
+            "url": url_prefix + attachment["path"],
+            "download_dir": attachment_path,
+            "attachment": attachment,
+        }
+        if entry is None:
+            urls.append(to_append)
         else:
-            shutil.copyfile(str(got_path), attachment_path / got_path.name)
+            try:
+                shutil.copyfile(str(entry["path"]), attachment_path / entry["name"])
+            except:
+                # If copy fails, the file either does not exist or we don't have
+                # permission. Download the file again and remove from db.
+                urls.append(to_append)
+                db.remove(entry["hash"])
 
     # Step attachments
     for step in testcase_info["steps"]:
         for attachment in step["attachments"]:
-            got_path = db.exists(attachment["hash"])
-            if got_path is None:
-                urls.append({
-                    "url": url_prefix + attachment["path"],
-                    "download_dir": attachment_path,
-                    "attachment": attachment,
-                })
+            entry = db.exists(attachment["hash"])
+            to_append = {
+                "url": url_prefix + attachment["path"],
+                "download_dir": attachment_path,
+                "attachment": attachment,
+            }
+            if entry is None:
+                urls.append(to_append)
             else:
-                shutil.copyfile(str(got_path), str(attachment_path / got_path.name))
+                try:
+                    shutil.copyfile(str(entry["path"]), attachment_path / entry["name"])
+                except:
+                    # If copy fails, the file either does not exist or we don't have
+                    # permission. Download the file again and remove from db.
+                    urls.append(to_append)
+                    db.remove(entry["hash"])
 
     results = ThreadPool(4).imap_unordered(download_attachment, urls)
     for r in results:
@@ -1430,7 +1444,7 @@ def download_attachments(testcase_info):
         if put:
             # If entry is successful, we copy the downloaded attachment to the
             # db directory.
-            shutil.copyfile(r["path"], str(attachment_path_in_db))
+            shutil.copyfile(r["path"], put["path"])
 
 
 # main function

Framework/attachment_db.py

@@ -4,16 +4,24 @@
 from pathlib import Path
 import time
 from typing import Any, Dict, Union
+import random
+import string
+
+
+def random_string(n: int) -> str:
+    return ''.join(random.choice(string.ascii_lowercase) for _ in range(n))
 
 
 class AttachmentDB:
     def __init__(self, db_directory: Path) -> None:
         self.db_directory = db_directory
         self.db_file = db_directory / "db.json"
         self.init_db()
+        self.path_suffix = ".zeuz."
+        self.suffix_length = 8
 
 
-    def exists(self, hash: str) -> Union[Path, None]:
+    def exists(self, hash: str) -> Union[Dict[str, str], None]:
         """
         exists returns a Path indicating whether the attachment exists in the
         database. None is returned if it does not exist.
@@ -25,23 +33,51 @@ def exists(self, hash: str) -> Union[Path, None]:
 
         if hash in db:
             entry = db[hash]
-            return Path(entry["path"])
+            return entry
 
         return None
 
 
+    def remove(self, hash: str) -> bool:
+        """
+        remove removes an attachment with the given hash from the db and returns
+        True if successful.
+        """
+
+        db = self.get_db()
+
+        if hash in db:
+            del db[hash]
+            self.save_db(db)
+            return True
+
+        return False
+
+
     def put(self, filepath: Path, hash: str):
         """
         put puts the attachment into the db.
         """
 
         if len(hash) == 0 or hash == "0":
-            return False
+            return None
 
         modified_at = time.time()
+
+        # We add a random suffix to the filepath to make sure files with same
+        # names but different hashes do not overwrite each other. Specially
+        # important if there are multiple attachments across multiple test
+        # cases/steps with the same file name.
+        path = filepath.with_suffix(
+            filepath.suffix +
+            self.path_suffix +
+            random_string(self.suffix_length)
+        )
+
         entry = {
             "hash": hash,
-            "path": str(filepath),
+            "path": str(path),
+            "name": filepath.name,
             "modified_at": modified_at,
         }
 
@@ -52,7 +88,7 @@ def put(self, filepath: Path, hash: str):
 
         self.save_db(db)
 
-        return True
+        return entry
 
 
     def get_db(self) -> Dict[str, Any]: