refactor: updated upscale_task creation to include background tasks

Author: JanssenBrm

app/database/models/processing_job.py

@@ -56,7 +56,7 @@ def save_job_to_db(
     db_session.add(job)
     db_session.commit()
     db_session.refresh(job)  # Refresh to get the ID after commit
-    logger.debug("Processing job saved with ID: {job.id}")
+    logger.debug(f"Processing job saved with ID: {job.id}")
     return job
 
 

app/routers/upscale_tasks.py

@@ -2,6 +2,7 @@
 import json
 from typing import Annotated
 from fastapi import (
+    BackgroundTasks,
     Body,
     APIRouter,
     Depends,
@@ -25,7 +26,11 @@
     UpscalingTaskSummary,
 )
 from app.schemas.websockets import WSTaskStatusMessage
-from app.services.upscaling import create_upscaling_task, get_upscaling_task_by_user_id
+from app.services.upscaling import (
+    create_upscaling_processing_jobs,
+    create_upscaling_task,
+    get_upscaling_task_by_user_id,
+)
 
 # from app.auth import get_current_user
 
@@ -98,12 +103,21 @@ async def create_upscale_task(
             },
         ),
     ],
+    background_tasks: BackgroundTasks,
     db: Session = Depends(get_db),
     user: str = "foobar",
 ) -> UpscalingTaskSummary:
     """Create a new upscaling job with the provided data."""
     try:
-        return create_upscaling_task(db, user, payload)
+        task = create_upscaling_task(db, user, payload)
+        background_tasks.add_task(
+            create_upscaling_processing_jobs,
+            database=db,
+            user=user,
+            request=payload,
+            upscaling_task_id=task.id,
+        )
+        return task
     except Exception as e:
         logger.exception(f"Error creating upscale task for user {user}: {e}")
         raise HTTPException(

app/services/upscaling.py

@@ -29,7 +29,7 @@
 }
 
 
-def _create_upscaling_processing_jobs(
+def create_upscaling_processing_jobs(
     database: Session, user: str, request: UpscalingTaskRequest, upscaling_task_id: int
 ) -> List[ProcessingJobSummary]:
     jobs: List[ProcessingJobSummary] = []
@@ -72,12 +72,6 @@ def create_upscaling_task(
         service=request.service.model_dump_json(),
     )
     record = save_upscaling_task_to_db(database, record)
-
-    logger.info(f"Creating upscaling job for {user} with request: {request}")
-    _create_upscaling_processing_jobs(
-        database=database, user=user, request=request, upscaling_task_id=record.id
-    )
-
     return UpscalingTaskSummary(
         id=record.id, title=record.title, label=record.label, status=record.status
     )

guides/upscaling_example.ipynb

@@ -81,3 +81,15 @@ def test_unit_jobs_get_job_results_404(mock_get_processing_job_results, client):
     r = client.get("/unit_jobs/1/results")
     assert r.status_code == 404
     assert "result for processing job 1 not found" in r.json().get("detail", "").lower()
+
+
+@patch("app.routers.unit_jobs.get_processing_job_results")
+def test_unit_jobs_get_job_results_500(mock_get_processing_job_results, client):
+
+    mock_get_processing_job_results.side_effect = RuntimeError(
+        "Database connection lost"
+    )
+
+    r = client.get("/unit_jobs/1/results")
+    assert r.status_code == 500
+    assert "database connection lost" in r.json().get("detail", "").lower()

tests/routers/test_unit_jobs.py

@@ -5,9 +5,11 @@
 import pytest
 
 
+@patch("app.routers.upscale_tasks.create_upscaling_processing_jobs")
 @patch("app.routers.upscale_tasks.create_upscaling_task")
 def test_upscaling_task_create_201(
     mock_create_upscaling_task,
+    mock_create_processing_jobs,
     client,
     fake_upscaling_task_request,
     fake_upscaling_task_summary,
@@ -18,6 +20,7 @@ def test_upscaling_task_create_201(
     r = client.post("/upscale_tasks", json=fake_upscaling_task_request.model_dump())
     assert r.status_code == 201
     assert r.json() == fake_upscaling_task_summary.model_dump()
+    assert mock_create_processing_jobs.called_once()
 
 
 @patch("app.routers.upscale_tasks.create_upscaling_task")
@@ -63,9 +66,7 @@ def test_upscaling_task_get_task_404(mock_get_upscale_task, client):
 
 
 @pytest.mark.asyncio
-@patch(
-    "app.routers.upscale_tasks.get_upscale_task", new_callable=AsyncMock
-)
+@patch("app.routers.upscale_tasks.get_upscale_task", new_callable=AsyncMock)
 async def test_ws_jobs_status(mock_get_task_status, client, fake_upscaling_task):
     mock_get_task_status.return_value = fake_upscaling_task
 
@@ -77,9 +78,7 @@ async def test_ws_jobs_status(mock_get_task_status, client, fake_upscaling_task)
 
 
 @pytest.mark.asyncio
-@patch(
-    "app.routers.upscale_tasks.get_upscale_task", new_callable=AsyncMock
-)
+@patch("app.routers.upscale_tasks.get_upscale_task", new_callable=AsyncMock)
 async def test_ws_jobs_status_closes_on_error(mock_get_task_status, client):
     mock_get_task_status.side_effect = RuntimeError("Database connection lost")
 
@@ -90,3 +89,18 @@ async def test_ws_jobs_status_closes_on_error(mock_get_task_status, client):
             websocket.receive_json()
 
         assert exc_info.value.code == 1011
+
+
+@pytest.mark.asyncio
+@patch("app.routers.upscale_tasks.get_upscale_task", new_callable=AsyncMock)
+async def test_ws_jobs_status_not_found(
+    mock_get_task_status, client, fake_upscaling_task
+):
+    mock_get_task_status.return_value = None
+
+    with client.websocket_connect("/ws/upscale_tasks/1?interval=1") as websocket:
+        websocket.receive_json()
+        websocket.receive_json()
+        data = websocket.receive_json()
+        assert data["type"] == "error"
+        assert data["message"].lower() == "task not found"

tests/routers/test_upscale_tasks.py

@@ -10,6 +10,7 @@
 from app.services.upscaling import (
     _get_upscale_status,
     _refresh_record_status,
+    create_upscaling_processing_jobs,
     create_upscaling_task,
     get_upscaling_task_by_user_id,
     get_upscaling_tasks_by_user_id,
@@ -38,21 +39,34 @@ def make_upscaling_record(status: ProcessingJobSummary) -> UpscalingTaskRecord:
     )
 
 
-@patch("app.services.upscaling.create_processing_job")
 @patch("app.services.upscaling.save_upscaling_task_to_db")
-def test_create_upscaling_task_creates_jobs(
+def test_create_upscaling_task_creates_task(
     mock_save_upscaling_task,
-    mock_create_processing_job,
     fake_upscaling_task_request,
     fake_upscaling_task_record,
     fake_upscaling_task_summary,
-    fake_processing_job_summary,
     fake_db_session,
 ):
     user = "foobar"
     mock_save_upscaling_task.return_value = fake_upscaling_task_record
-    mock_create_processing_job.return_value = fake_processing_job_summary
     result = create_upscaling_task(fake_db_session, user, fake_upscaling_task_request)
+    mock_save_upscaling_task.assert_called_once()
+    assert result == fake_upscaling_task_summary
+
+
+@patch("app.services.upscaling.create_processing_job")
+def test_create_upscaling_task_creates_jobs(
+    mock_create_processing_job,
+    fake_upscaling_task_request,
+    fake_upscaling_task_record,
+    fake_processing_job_summary,
+    fake_db_session,
+):
+    user = "foobar"
+    mock_create_processing_job.return_value = fake_processing_job_summary
+    result = create_upscaling_processing_jobs(
+        fake_db_session, user, fake_upscaling_task_request, 1
+    )
 
     expected_calls = [
         call(
@@ -73,8 +87,7 @@ def test_create_upscaling_task_creates_jobs(
     ]
 
     mock_create_processing_job.assert_has_calls(expected_calls)
-    mock_save_upscaling_task.assert_called_once()
-    assert result == fake_upscaling_task_summary
+    assert len(result) == len(fake_upscaling_task_request.dimension.values)
 
 
 def test_returns_running_if_any_running():