Fix temporary file cleanup for downloads/zips

mcserver/tasks.py

@@ -24,7 +24,8 @@
 from mcserver.zipsession_v2 import (
     SessionDirectoryConstructor,
     SubjectDirectoryConstructor,
-    zipdir
+    zipdir,
+    rmtree_with_retry
 )
 
 
@@ -51,57 +52,46 @@ def download_session_archive(self, session_id, user_id=None):
     """ This task is responsible for asynchronous session archive download.
         If user_id is None, the public session download occurred.
     """
-    import shutil
-
-    session_dir_path = None
+    # Known up front so the build dir is cleaned up even if build() raises.
+    session_dir_path = os.path.join(
+        settings.MEDIA_ROOT, f"OpenCapData_{session_id}")
     session_zip_path = None
 
     try:
         session_dir_path = SessionDirectoryConstructor().build(session_id)
         session_zip_path = zipdir(session_dir_path)
-
         create_download_log(session_zip_path, self.request.id, user_id)
-        os_remove_with_retry(session_zip_path)
-
     except Exception as e:
-        # Delete files and send the traceback to Sentry if something went wrong
-        if session_dir_path and os.path.isfile(session_dir_path):
-            shutil.rmtree(session_dir_path)
-        if session_zip_path and os.path.isfile(session_zip_path):
-            os.remove(session_zip_path)
         if settings.SENTRY_DSN:
             import sentry_sdk
             sentry_sdk.capture_exception(e)
         else:
             print(e)
+    finally:
+        cleanup_download_tmp_files(session_dir_path, session_zip_path)
+
 
 @shared_task(bind=True)
 def download_subject_archive(self, subject_id, user_id):
     """ This task is responsible for asynchronous subject archive download
     """
-    import shutil
-
-    subject_dir_path = None
+    # Known up front so the build dir is cleaned up even if build() raises.
+    subject_dir_path = os.path.join(
+        settings.MEDIA_ROOT, f"OpenCapData_Subject_{subject_id}")
     subject_zip_path = None
 
     try:
         subject_dir_path = SubjectDirectoryConstructor().build(subject_id)
         subject_zip_path = zipdir(subject_dir_path)
-        with open(subject_zip_path, "rb") as archive:
-            log = DownloadLog.objects.create(task_id=str(self.request.id), user_id=user_id)
-            log.media.save(os.path.basename(subject_zip_path), archive)
-            os.remove(subject_zip_path)
+        create_download_log(subject_zip_path, self.request.id, user_id)
     except Exception as e:
-        # Delete files and send the traceback to Sentry if something went wrong
-        if subject_dir_path and os.path.isfile(subject_dir_path):
-            shutil.rmtree(subject_dir_path)
-        if subject_zip_path and os.path.isfile(subject_zip_path):
-            os.remove(subject_zip_path)
         if settings.SENTRY_DSN:
             import sentry_sdk
             sentry_sdk.capture_exception(e)
         else:
             print(e)
+    finally:
+        cleanup_download_tmp_files(subject_dir_path, subject_zip_path)
 
 
 @shared_task
@@ -199,7 +189,15 @@ def submit_cloudwatch_metrics():
     submit_number_of_pending_trials_to_cloudwatch()
 
 # Helper functions
-def create_download_log(zip_path, task_id, user_id, 
+def cleanup_download_tmp_files(dir_path, zip_path):
+    """Remove the local build dir and zip archive for a download, if present."""
+    if dir_path and os.path.isdir(dir_path):
+        rmtree_with_retry(dir_path)
+    if zip_path and os.path.isfile(zip_path):
+        os_remove_with_retry(zip_path)
+
+
+def create_download_log(zip_path, task_id, user_id,
                         max_retries=5, backoff=0.1):
     archive = None
     for attempt in range(max_retries):

tests/test_tasks.py

@@ -1,8 +1,11 @@
 import os
 import json
+import shutil
 import tempfile
+import zipfile
 from unittest import mock
 
+from django.conf import settings
 from django.test import TestCase, override_settings
 
 from mcserver.models import (
@@ -27,6 +30,7 @@
 )
 
 _temp_media = tempfile.mkdtemp()
+_download_tmp = tempfile.mkdtemp()  # isolated dir for DownloadArchiveTests' tearDown
 
 
 @override_settings(
@@ -54,66 +58,6 @@ def setUp(self):
         self.trial_one = Trial.objects.create(session=self.session, name="testone")
         self.trial_two = Trial.objects.create(session=self.session, name="testtwo")
 
-    @mock.patch("mcserver.tasks.zipdir")
-    @mock.patch.object(SessionDirectoryConstructor, "build")
-    def test_download_session_archive_creates_archive_and_logs_action(
-        self, mock_dir_builder, mock_zipdir
-    ):
-        mock_zipdir.return_value = "archive.zip"
-        mock_dir_builder.return_value = "archive"
-        before_logs = DownloadLog.objects.count()
-        task = download_session_archive.delay("dummy-session-id", self.user.id)
-        after_logs = DownloadLog.objects.count()
-        self.assertEqual(after_logs, before_logs + 1)
-
-        log = DownloadLog.objects.last()
-        self.assertEqual(log.task_id, task.id)
-        self.assertEqual(log.user, self.user)
-        self.assertEqual(log.media_path, "archive.zip")
-
-        mock_dir_builder.assert_called_once_with("dummy-session-id")
-        mock_zipdir.assert_called_once_with("archive")
-
-    @mock.patch("mcserver.tasks.zipdir")
-    @mock.patch.object(SessionDirectoryConstructor, "build")
-    def test_download_session_archive_creates_archive_and_logs_action_for_anon_user(
-        self, mock_dir_builder, mock_zipdir
-    ):
-        mock_zipdir.return_value = "archive.zip"
-        mock_dir_builder.return_value = "archive"
-        before_logs = DownloadLog.objects.count()
-        task = download_session_archive.delay("dummy-session-id", None)
-        after_logs = DownloadLog.objects.count()
-        self.assertEqual(after_logs, before_logs + 1)
-
-        log = DownloadLog.objects.last()
-        self.assertEqual(log.task_id, task.id)
-        self.assertIsNone(log.user)
-        self.assertEqual(log.media_path, "archive.zip")
-
-        mock_dir_builder.assert_called_once_with("dummy-session-id")
-        mock_zipdir.assert_called_once_with("archive")
-
-    @mock.patch("mcserver.tasks.zipdir")
-    @mock.patch.object(SubjectDirectoryConstructor, "build")
-    def test_download_subject_archive_creates_archive_and_logs_action(
-        self, mock_dir_builder, mock_zipdir
-    ):
-        mock_zipdir.return_value = "archive.zip"
-        mock_dir_builder.return_value = "archive"
-        before_logs = DownloadLog.objects.count()
-        task = download_subject_archive.delay("dummy-subject-id", self.user.id)
-        after_logs = DownloadLog.objects.count()
-        self.assertEqual(after_logs, before_logs + 1)
-
-        log = DownloadLog.objects.last()
-        self.assertEqual(log.task_id, task.id)
-        self.assertEqual(log.user, self.user)
-        self.assertEqual(log.media_path, "archive.zip")
-
-        mock_dir_builder.assert_called_once_with("dummy-subject-id")
-        mock_zipdir.assert_called_once_with("archive")
-
     def test_delete_pingdom_sessions_successful(self):
         Session.objects.create(user=self.pingdom_user)
         Session.objects.create(user=self.pingdom_user)
@@ -137,7 +81,7 @@ def test_delete_pingdom_sessions_no_sessions(self):
         self.assertFalse(Session.objects.filter(user=self.pingdom_user).exists())
         delete_pingdom_sessions.delay()
         self.assertFalse(Session.objects.filter(user=self.pingdom_user).exists())
-    
+
     @mock.patch("requests.post")
     def test_invoke_aws_lambda_function_commits_successful_analysis_result(
         self, mock_post_request
@@ -201,7 +145,7 @@ def test_invoke_aws_lambda_function_commits_failed_analysis_result_if_aws_error(
         self.assertIsNone(analysis_result.result)
         self.assertEqual(analysis_result.response, {'error': 'session_id is required.'})
         self.assertEqual(analysis_result.state, AnalysisResultState.FAILED)
-    
+
     def test_invoke_aws_lambda_function_commits_failed_analysis_result_if_request_exception(
         self
     ):
@@ -225,7 +169,7 @@ def test_invoke_aws_lambda_function_commits_failed_analysis_result_if_request_ex
         )
         self.assertIsNone(analysis_result.result)
         self.assertEqual(analysis_result.state, AnalysisResultState.FAILED)
-    
+
     @mock.patch("requests.post")
     def test_invoke_aws_lambda_function_commits_failed_analysis_result_if_json_invalid(
         self, mock_post_request
@@ -252,7 +196,7 @@ def test_invoke_aws_lambda_function_commits_failed_analysis_result_if_json_inval
         self.assertIsNone(analysis_result.result)
         self.assertEqual(analysis_result.response, {'error': 'Invalid JSON.'})
         self.assertEqual(analysis_result.state, AnalysisResultState.FAILED)
-    
+
     @mock.patch("requests.post")
     def test_invoke_aws_lambda_function_re_run_analysis_function_in_the_same_result_instance(
         self, mock_post_request
@@ -285,4 +229,144 @@ def test_invoke_aws_lambda_function_re_run_analysis_function_in_the_same_result_
         self.assertEqual(analisys_result.data, data)
         self.assertEqual(analisys_result.status, status_code)
         self.assertEqual(analisys_result.result, result)
-        self.assertEqual(analisys_result.state, AnalysisResultState.SUCCESSFULL)
+        self.assertEqual(analisys_result.state, AnalysisResultState.SUCCESSFULL)
+
+
+@override_settings(
+    MEDIA_ROOT=_download_tmp,
+    ARCHIVES_ROOT=os.path.join(_download_tmp, "archives"),
+    DEFAULT_FILE_STORAGE="django.core.files.storage.FileSystemStorage",
+    SENTRY_DSN="",
+)
+class DownloadArchiveTests(TestCase):
+    """download_session_archive / download_subject_archive must create a
+    DownloadLog on success and never leave the build dir or zip on the
+    worker's ephemeral disk (which otherwise fills up and breaks downloads).
+    """
+
+    def setUp(self):
+        self.user = User.objects.create_user(username="dl-user", password="pw")
+        os.makedirs(settings.MEDIA_ROOT, exist_ok=True)
+
+    def tearDown(self):
+        for name in os.listdir(settings.MEDIA_ROOT):
+            path = os.path.join(settings.MEDIA_ROOT, name)
+            if os.path.isdir(path):
+                shutil.rmtree(path, ignore_errors=True)
+            else:
+                os.remove(path)
+
+    @staticmethod
+    def _make_fake_build(dir_path):
+        """Return a fake build() (Session/SubjectDirectoryConstructor.build)
+        that writes a real build dir on disk and returns its path."""
+        def _build(object_id):
+            os.makedirs(dir_path, exist_ok=True)
+            with open(os.path.join(dir_path, "payload.txt"), "w") as fh:
+                fh.write("data")
+            return dir_path
+        return _build
+
+    @staticmethod
+    def _fake_zipdir(dir_path):
+        """Fake zipdir() for the success path, mirroring the real one: remove
+        the source dir and write a real zip under ARCHIVES_ROOT."""
+        shutil.rmtree(dir_path)
+        os.makedirs(settings.ARCHIVES_ROOT, exist_ok=True)
+        zip_path = os.path.join(
+            settings.ARCHIVES_ROOT, os.path.basename(dir_path) + ".zip")
+        with zipfile.ZipFile(zip_path, "w") as zf:
+            zf.writestr("payload.txt", "data")
+        return zip_path
+
+    # --- success: DownloadLog created and temp files cleaned up ---------
+
+    def test_session_success_creates_log_and_cleans_up(self):
+        build_dir = os.path.join(settings.MEDIA_ROOT, "OpenCapData_sess-1")
+        leftover_zip = os.path.join(
+            settings.ARCHIVES_ROOT, "OpenCapData_sess-1.zip")
+        with mock.patch.object(
+            SessionDirectoryConstructor, "build",
+            side_effect=self._make_fake_build(build_dir),
+        ), mock.patch(
+            "mcserver.tasks.zipdir", side_effect=self._fake_zipdir,
+        ):
+            result = download_session_archive.apply(args=("sess-1", self.user.id))
+
+        self.assertEqual(DownloadLog.objects.count(), 1)
+        log = DownloadLog.objects.get()
+        self.assertEqual(log.user, self.user)
+        self.assertEqual(log.task_id, result.id)
+        self.assertTrue(log.media.name)  # a file was actually stored
+        self.assertFalse(os.path.exists(build_dir))
+        self.assertFalse(os.path.exists(leftover_zip))
+
+    def test_session_success_for_anonymous_user(self):
+        build_dir = os.path.join(settings.MEDIA_ROOT, "OpenCapData_sess-anon")
+        with mock.patch.object(
+            SessionDirectoryConstructor, "build",
+            side_effect=self._make_fake_build(build_dir),
+        ), mock.patch(
+            "mcserver.tasks.zipdir", side_effect=self._fake_zipdir,
+        ):
+            download_session_archive.apply(args=("sess-anon", None))
+
+        self.assertEqual(DownloadLog.objects.count(), 1)
+        self.assertIsNone(DownloadLog.objects.get().user)
+
+    def test_subject_success_creates_log_and_cleans_up(self):
+        build_dir = os.path.join(
+            settings.MEDIA_ROOT, "OpenCapData_Subject_subj-1")
+        leftover_zip = os.path.join(
+            settings.ARCHIVES_ROOT, "OpenCapData_Subject_subj-1.zip")
+        with mock.patch.object(
+            SubjectDirectoryConstructor, "build",
+            side_effect=self._make_fake_build(build_dir),
+        ), mock.patch(
+            "mcserver.tasks.zipdir", side_effect=self._fake_zipdir,
+        ):
+            result = download_subject_archive.apply(args=("subj-1", self.user.id))
+
+        self.assertEqual(DownloadLog.objects.count(), 1)
+        log = DownloadLog.objects.get()
+        self.assertEqual(log.user, self.user)
+        self.assertEqual(log.task_id, result.id)
+        self.assertFalse(os.path.exists(build_dir))
+        self.assertFalse(os.path.exists(leftover_zip))
+
+    # --- failure: no DownloadLog, but temp still cleaned up ------------
+
+    def test_session_cleans_up_build_dir_when_zip_fails(self):
+        build_dir = os.path.join(settings.MEDIA_ROOT, "OpenCapData_sess-2")
+        with mock.patch.object(
+            SessionDirectoryConstructor, "build",
+            side_effect=self._make_fake_build(build_dir),
+        ), mock.patch(
+            "mcserver.tasks.zipdir",
+            side_effect=OSError("No space left on device"),
+        ):
+            download_session_archive.apply(args=("sess-2", self.user.id))
+
+        self.assertFalse(
+            os.path.exists(build_dir),
+            "build dir must be removed after a failed download",
+        )
+        self.assertEqual(DownloadLog.objects.count(), 0)
+
+    def test_subject_cleans_up_build_dir_when_zip_fails(self):
+        build_dir = os.path.join(
+            settings.MEDIA_ROOT, "OpenCapData_Subject_subj-2")
+        with mock.patch.object(
+            SubjectDirectoryConstructor, "build",
+            side_effect=self._make_fake_build(build_dir),
+        ), mock.patch(
+            "mcserver.tasks.zipdir",
+            side_effect=OSError("No space left on device"),
+        ):
+            download_subject_archive.apply(args=("subj-2", self.user.id))
+
+        self.assertFalse(
+            os.path.exists(build_dir),
+            "build dir must be removed after a failed download",
+        )
+        self.assertEqual(DownloadLog.objects.count(), 0)