diff --git a/mcserver/tasks.py b/mcserver/tasks.py index 6f0391d..3baceff 100644 --- a/mcserver/tasks.py +++ b/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): diff --git a/tests/test_tasks.py b/tests/test_tasks.py index 1ebabb1..b08212e 100644 --- a/tests/test_tasks.py +++ b/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) \ No newline at end of file + 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)