Merge ~ines-almeida/launchpad:prevent-stray-temp-file-on-exception into launchpad:master

[Ines Almeida]

Cleanup temp files if copying component raises

If something goes wrong while copying a component, we:
 * catch the exception (generally, whichever exception)
 * do some cleanup of temporary files
 * re-raise the same exception

LP: #1662671

Added a couple of unit tests that test:
 - the cleanup method removes the temporary path
 - the cleanup is called when an exception is raised, and the exact same exception is re-raised

[Colin Watson (cjwatson)]

review: Approve

diff --git a/lib/lp/archivepublisher/diskpool.py b/lib/lp/archivepublisher/diskpool.py
index c7b2576..18e4124 100644
--- a/lib/lp/archivepublisher/diskpool.py
+++ b/lib/lp/archivepublisher/diskpool.py
@@ -149,6 +149,11 @@ class _diskpool_atomicfile:
         # different filesystems.
         self.tempname.rename(self.targetfilename)
 
+    def cleanup_temporary_path(self) -> None:
+        """Removes temporary path created on __init__"""
+        if self.tempname.exists():
+            self.tempname.unlink()
+
 
 class DiskPoolEntry:
     """Represents a single file in the pool, across all components.
@@ -285,8 +290,16 @@ class DiskPoolEntry:
         file_to_write = _diskpool_atomicfile(
             targetpath, "wb", rootpath=self.temppath
         )
-        lfa.open()
-        copy_and_close(lfa, file_to_write)
+        try:
+            lfa.open()
+            copy_and_close(lfa, file_to_write)
+        except Exception:
+            # Prevent ending up with a stray temporary file lying around if
+            # anything goes wrong while copying the file. Still raises error
+            self.debug("Cleaning up temp path %s" % file_to_write.tempname)
+            file_to_write.cleanup_temporary_path()
+            raise
+
         self.file_component = component
         return FileAddActionEnum.FILE_ADDED
 
diff --git a/lib/lp/archivepublisher/tests/test_pool.py b/lib/lp/archivepublisher/tests/test_pool.py
index 311f203..1329e34 100644
--- a/lib/lp/archivepublisher/tests/test_pool.py
+++ b/lib/lp/archivepublisher/tests/test_pool.py
@@ -5,11 +5,17 @@
 
 import hashlib
 from pathlib import Path, PurePath
+from unittest import mock
 
 from lazr.enum import EnumeratedType, Item
 from zope.interface import alsoProvides, implementer
 
-from lp.archivepublisher.diskpool import DiskPool, poolify, unpoolify
+from lp.archivepublisher.diskpool import (
+    DiskPool,
+    _diskpool_atomicfile,
+    poolify,
+    unpoolify,
+)
 from lp.registry.interfaces.sourcepackage import SourcePackageFileType
 from lp.services.log.logger import BufferLogger
 from lp.soyuz.enums import ArchiveRepositoryFormat, BinaryPackageFileType
@@ -102,6 +108,10 @@ class FakePackageReleaseFile:
             alsoProvides(self, IBinaryPackageFile)
 
 
+class SpecificTestException(Exception):
+    pass
+
+
 class PoolTestingFile:
     def __init__(
         self,
@@ -290,3 +300,32 @@ class TestPool(TestCase):
         foo.removeFromPool("main")
         self.assertFalse(foo.checkExists("main"))
         self.assertTrue(foo.checkIsFile("universe"))
+
+    def test_raise_deletes_temporary_file(self):
+        """If copying fails, cleanup is called and same error is raised"""
+        foo = PoolTestingFile(
+            pool=self.pool,
+            source_name="foo",
+            source_version="1.0",
+            filename="foo-1.0.deb",
+        )
+
+        with mock.patch(
+            "lp.archivepublisher.diskpool.copy_and_close",
+            side_effect=SpecificTestException,
+        ), mock.patch.object(
+            _diskpool_atomicfile, "cleanup_temporary_path"
+        ) as mock_cleanup:
+            self.assertRaises(SpecificTestException, foo.addToPool, "universe")
+
+        self.assertEqual(mock_cleanup.call_count, 1)
+        self.assertFalse(foo.checkIsFile("universe"))
+
+    def test_diskpool_atomicfile(self):
+        """Temporary files are properly removed"""
+        target_path = Path(self.makeTemporaryDirectory() + "/temp.file")
+        foo = _diskpool_atomicfile(target_path, "wb")
+
+        self.assertTrue(foo.tempname.exists())
+        foo.cleanup_temporary_path()
+        self.assertFalse(foo.tempname.exists())