Add initial model uploads

Author: gfrn

src/pato/crud/sessions.py

@@ -1,3 +1,4 @@
+import os
 import pathlib
 import shutil
 from datetime import datetime
@@ -18,6 +19,7 @@
 from ..utils.generic import ProposalReference, check_session_active, parse_proposal
 
 HDF5_FILE_SIGNATURE = b"\x89\x48\x44\x46\x0d\x0a\x1a\x0a"
+MRC_FILE_SIGNATURE = b"\x4d\x41\x50"
 
 
 def _validate_session_active(proposal_reference: ProposalReference):
@@ -263,3 +265,43 @@ def upload_processing_model(file: UploadFile, proposal_reference: ProposalRefere
         )
 
     file.file.close()
+
+
+def upload_initial_model(file: UploadFile, proposal_reference: ProposalReference):
+    file.file.seek(208)
+    file_header = file.file.read(3)
+    file.file.seek(0)
+
+    if file_header != MRC_FILE_SIGNATURE:
+        raise HTTPException(
+            detail="Invalid file type (must be MRC file)",
+            status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
+        )
+
+    file_path = f"{_get_folder_and_visit(proposal_reference)[0]}/processing"
+
+    if not os.path.isdir(file_path):
+        app_logger.error(
+            f"Failed to upload {file.filename}: visit directory '{file_path}' does not exist"
+        )
+        raise HTTPException(
+            detail="Failed to upload file",
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+    file_path = os.path.join(file_path, "initial_model")
+    # Create initial_model folder if it does not yet exist
+    os.makedirs(file_path, exist_ok=True)
+
+    try:
+        with open(os.path.join(file_path, file.filename or "model.mrc"), "wb") as f:
+            shutil.copyfileobj(file.file, f)
+    except OSError as e:
+        file.file.close()
+        app_logger.error(f"Failed to upload {file.filename}: {e}")
+        raise HTTPException(
+            detail="Failed to upload file",
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        )
+
+    file.file.close()

src/pato/routes/proposals.py

@@ -20,6 +20,7 @@
     prefix="/proposals",
 )
 
+
 @router.get(
     "/{proposalReference}/sessions/{visitNumber}/dataGroups",
     response_model=Paged[DataCollectionGroupSummaryResponse],
@@ -93,3 +94,13 @@ def upload_processing_model(
     return sessions_crud.upload_processing_model(
         file=file, proposal_reference=proposalReference
     )
+
+
+@router.post("/{proposalReference}/sessions/{visitNumber}/initialModel")
+def upload_initial_model(
+    file: UploadFile, proposalReference=Depends(Permissions.session)
+):
+    """Upload custom initial model"""
+    return sessions_crud.upload_initial_model(
+        file=file, proposal_reference=proposalReference
+    )

tests/sessions/test_upload_initial_model.py

@@ -0,0 +1,64 @@
+from datetime import datetime
+from unittest.mock import mock_open, patch
+
+from lims_utils.tables import BLSession
+
+VALID_FILE = b"\x00" * 208 + b"\x4d\x41\x50"
+
+
+def active_mock(_):
+    return BLSession(
+        startDate=datetime(year=2022, month=1, day=1),
+        sessionId=27464088,
+        beamLineName="m12",
+    )
+
+
+@patch("pato.crud.sessions._validate_session_active", new=active_mock)
+@patch("builtins.open", new_callable=mock_open())
+@patch("pato.crud.sessions.os.path.isdir", new=lambda _: True)
+def test_post(_, mock_permissions, client):
+    """Should write file successfully if file matches expected signature"""
+    resp = client.post(
+        "/proposals/cm31111/sessions/5/initialModel",
+        files={"file": ("mrc.mrc", VALID_FILE, "application/octet-stream")},
+    )
+
+    assert resp.status_code == 200
+
+
+@patch("pato.crud.sessions._validate_session_active", new=active_mock)
+@patch("builtins.open", new_callable=mock_open())
+def test_invalid_file_signature(_, mock_permissions, client):
+    """Should raise exception if file signature doesn't match MRC file signature"""
+    resp = client.post(
+        "/proposals/cm31111/sessions/5/initialModel",
+        files={"file": ("not-mrc.mrc", b"\x01\x02", "application/octet-stream")},
+    )
+
+    assert resp.status_code == 415
+
+
+@patch("pato.crud.sessions._validate_session_active", new=active_mock)
+@patch("builtins.open", side_effect=OSError("Write Error"))
+def test_write_error(_, mock_permissions, client):
+    """Should return 500 if there was an error writing the file"""
+    resp = client.post(
+        "/proposals/cm31111/sessions/5/initialModel",
+        files={"file": ("mrc.mrc", VALID_FILE, "application/octet-stream")},
+    )
+
+    assert resp.status_code == 500
+
+
+@patch("pato.crud.sessions._validate_session_active", new=active_mock)
+@patch("builtins.open", side_effect=OSError("Write Error"))
+@patch("pato.crud.sessions.os.path.isdir", new=lambda _: False)
+def test_dir_does_not_exist(_, mock_permissions, client):
+    """Should return 500 if directory does not exist"""
+    resp = client.post(
+        "/proposals/cm31111/sessions/5/initialModel",
+        files={"file": ("mrc.mrc", VALID_FILE, "application/octet-stream")},
+    )
+
+    assert resp.status_code == 500