From 8235dcc4f7dab131dc268a8d761d98977c877fd5 Mon Sep 17 00:00:00 2001
From: Mike Cantelon <mcantelon@gmail.com>
Date: Mon, 16 Jan 2023 18:41:22 -0800
Subject: [PATCH] Add Celery task for file cleanup after fetch job

Adds a small Celery task to delete the directory in which files related
to a fetch job are stored.
---
 AIPscan/Aggregator/tasks.py            | 10 ++++++++++
 AIPscan/Aggregator/tests/test_tasks.py | 17 +++++++++++++++++
 AIPscan/Aggregator/views.py            |  3 +++
 3 files changed, 30 insertions(+)

diff --git a/AIPscan/Aggregator/tasks.py b/AIPscan/Aggregator/tasks.py
index ff46ada2..8eebdc63 100644
--- a/AIPscan/Aggregator/tasks.py
+++ b/AIPscan/Aggregator/tasks.py
@@ -281,6 +281,16 @@ def package_lists_request(self, storage_service_id, timestamp, packages_director
     }
 
 
+@celery.task()
+def fetch_job_file_cleanup(fetch_job_id):
+    obj = FetchJob.query.filter_by(id=fetch_job_id).first()
+
+    if os.path.isdir(obj.download_directory):
+        shutil.rmtree(obj.download_directory)
+
+    logger.info("Cleaned up after fetch job {}".format(fetch_job_id))
+
+
 @celery.task()
 def get_mets(
     package_uuid,
diff --git a/AIPscan/Aggregator/tests/test_tasks.py b/AIPscan/Aggregator/tests/test_tasks.py
index b737aeba..6b09d27a 100644
--- a/AIPscan/Aggregator/tests/test_tasks.py
+++ b/AIPscan/Aggregator/tests/test_tasks.py
@@ -13,6 +13,7 @@
     delete_aip,
     delete_fetch_job,
     delete_storage_service,
+    fetch_job_file_cleanup,
     get_mets,
     make_request,
     parse_package_list_file,
@@ -212,6 +213,22 @@ def test_delete_storage_service_task(app_instance, tmpdir, mocker):
     assert Agent.query.filter_by(storage_service_id=storage_service.id).count() == 0
 
 
+def test_fetch_job_file_cleanup_task(app_instance, tmpdir, mocker):
+    test_downloads_dir = os.path.join(tmpdir, "downloads")
+    os.mkdir(test_downloads_dir)
+
+    # Mock download directory should be created
+    assert os.path.isdir(test_downloads_dir)
+
+    fetch_job1 = test_helpers.create_test_fetch_job(
+        download_directory=test_downloads_dir
+    )
+    fetch_job_file_cleanup(fetch_job1.id)
+
+    # Mock download directory should be removed
+    assert not os.path.isdir(test_downloads_dir)
+
+
 @pytest.mark.parametrize(
     "response, raises_task_error",
     [
diff --git a/AIPscan/Aggregator/views.py b/AIPscan/Aggregator/views.py
index 60ba7770..fd3cceed 100644
--- a/AIPscan/Aggregator/views.py
+++ b/AIPscan/Aggregator/views.py
@@ -287,6 +287,9 @@ def get_mets_task_status(coordinatorid):
     downloadStart = _format_date(start)
     obj.download_end = downloadEnd
     db.session.commit()
+
+    tasks.fetch_job_file_cleanup.delay(fetchJobId)
+
     response = {"state": "COMPLETED"}
     flash("Fetch Job {} completed".format(downloadStart))
     return jsonify(response)