Use APScheduler

Author: BenjaminCharmes

pydatalab/src/pydatalab/routes/v0_1/export.py

@@ -21,48 +21,48 @@
 def _(): ...
 
 
+def _do_export(task_id: str, collection_id: str):
+    try:
+        flask_mongo.db.export_tasks.update_one(
+            {"task_id": task_id}, {"$set": {"status": ExportStatus.PROCESSING}}
+        )
+
+        export_dir = Path(CONFIG.FILE_DIRECTORY) / "exports"
+        export_dir.mkdir(exist_ok=True)
+
+        output_path = export_dir / f"{task_id}.eln"
+        create_eln_file(collection_id, str(output_path))
+
+        flask_mongo.db.export_tasks.update_one(
+            {"task_id": task_id},
+            {
+                "$set": {
+                    "status": ExportStatus.READY,
+                    "file_path": str(output_path),
+                    "completed_at": datetime.now(tz=timezone.utc),
+                }
+            },
+        )
+
+    except Exception as e:
+        flask_mongo.db.export_tasks.update_one(
+            {"task_id": task_id},
+            {
+                "$set": {
+                    "status": ExportStatus.ERROR,
+                    "error_message": str(e),
+                    "completed_at": datetime.now(tz=timezone.utc),
+                }
+            },
+        )
+
+
 def _generate_export_in_background(task_id: str, collection_id: str, app):
-    """Background function to generate the .eln file.
-
-    Parameters:
-        task_id: ID of the export task
-        collection_id: ID of the collection to export
-        app: Flask application instance
-    """
-    with app.app_context():
-        try:
-            flask_mongo.db.export_tasks.update_one(
-                {"task_id": task_id}, {"$set": {"status": ExportStatus.PROCESSING}}
-            )
-
-            export_dir = Path(CONFIG.FILE_DIRECTORY) / "exports"
-            export_dir.mkdir(exist_ok=True)
-
-            output_path = export_dir / f"{task_id}.eln"
-            create_eln_file(collection_id, str(output_path))
-
-            flask_mongo.db.export_tasks.update_one(
-                {"task_id": task_id},
-                {
-                    "$set": {
-                        "status": ExportStatus.READY,
-                        "file_path": str(output_path),
-                        "completed_at": datetime.now(tz=timezone.utc),
-                    }
-                },
-            )
-
-        except Exception as e:
-            flask_mongo.db.export_tasks.update_one(
-                {"task_id": task_id},
-                {
-                    "$set": {
-                        "status": ExportStatus.ERROR,
-                        "error_message": str(e),
-                        "completed_at": datetime.now(tz=timezone.utc),
-                    }
-                },
-            )
+    if app is not None:
+        with app.app_context():
+            _do_export(task_id, collection_id)
+    else:
+        _do_export(task_id, collection_id)
 
 
 @EXPORT.route("/collections/<string:collection_id>/export", methods=["POST"])
@@ -82,7 +82,7 @@ def start_collection_export(collection_id: str):
     task_id = str(uuid.uuid4())
 
     if not CONFIG.TESTING:
-        creator_id = current_user.person.immutable_id
+        creator_id = str(current_user.person.immutable_id)
     else:
         creator_id = "000000000000000000000000"
 
@@ -95,7 +95,10 @@ def start_collection_export(collection_id: str):
 
     flask_mongo.db.export_tasks.insert_one(export_task.dict())
 
-    app = current_app._get_current_object()
+    try:
+        app = current_app._get_current_object()
+    except RuntimeError:
+        app = None
 
     export_scheduler.add_job(
         func=_generate_export_in_background,
@@ -110,15 +113,6 @@ def start_collection_export(collection_id: str):
 
 @EXPORT.route("/exports/<string:task_id>/status", methods=["GET"])
 def get_export_status(task_id: str):
-    """Get the status of an export task.
-
-    Parameters:
-        task_id: The export task ID
-
-    Returns:
-        JSON response with task status and download URL if ready
-    """
-
     task = flask_mongo.db.export_tasks.find_one({"task_id": task_id})
 
     if not task:
@@ -146,15 +140,6 @@ def get_export_status(task_id: str):
 
 @EXPORT.route("/exports/<string:task_id>/download", methods=["GET"])
 def download_export(task_id: str):
-    """Download the generated .eln file.
-
-    Parameters:
-        task_id: The export task ID
-
-    Returns:
-        The .eln file as attachment
-    """
-
     task = flask_mongo.db.export_tasks.find_one({"task_id": task_id})
 
     if not task:

pydatalab/tests/server/test_export.py

@@ -27,35 +27,37 @@ def sample_collection(database, user_id):
 
 @pytest.fixture
 def mock_scheduler():
-    with patch("pydatalab.routes.v0_1.export.export_scheduler") as mock_scheduler:
-        mock_scheduler.add_job = MagicMock()
-
-        def mock_add_job(func, args, job_id=None):
-            return None
-
-        mock_scheduler.add_job.side_effect = mock_add_job
-
-        yield mock_scheduler
+    with patch("pydatalab.routes.v0_1.export.export_scheduler") as mock_sched:
+        mock_sched.add_job = MagicMock(return_value=None)
+        yield mock_sched
 
 
 class TestExportRoutes:
-    def test_start_collection_export_success(self, client, sample_collection, mock_scheduler):
+    def test_start_collection_export_success(
+        self, client, sample_collection, mock_scheduler, database
+    ):
         collection_id = sample_collection["collection_id"]
 
-        with patch("pydatalab.routes.v0_1.export.current_app") as mock_app:
-            mock_app._get_current_object.return_value = MagicMock()
+        response = client.post(f"/collections/{collection_id}/export")
 
-            response = client.post(f"/collections/{collection_id}/export")
+        assert response.status_code == 202
+
+        data = json.loads(response.data)
+        assert data["status"] == "success"
+        assert "task_id" in data
+        assert "status_url" in data
+        assert data["status_url"] == f"/exports/{data['task_id']}/status"
 
-            assert response.status_code == 202
+        assert mock_scheduler.add_job.called
 
-            data = json.loads(response.data)
-            assert data["status"] == "success"
-            assert "task_id" in data
-            assert "status_url" in data
-            assert data["status_url"] == f"/exports/{data['task_id']}/status"
+        # Verify the task was created in the database
+        task = database.export_tasks.find_one({"task_id": data["task_id"]})
+        assert task is not None
+        assert task["collection_id"] == collection_id
+        assert task["status"] == ExportStatus.PENDING
 
-            assert mock_scheduler.add_job.called
+        # Clean up
+        database.export_tasks.delete_one({"task_id": data["task_id"]})
 
     def test_start_collection_export_not_found(self, client):
         response = client.post("/collections/nonexistent/export")