Merge pull request #25 from The-Strategy-Unit/v5-fixes

v5 fixes

Author: tomjemmett

src/nhp/aci/run_model/storage.py

@@ -1,6 +1,7 @@
 """Methods for uploading model parameters to Azure Blob Storage."""
 
 import logging
+from datetime import datetime, timezone
 from typing import Any
 
 from azure.core.credentials import TokenCredential
@@ -67,4 +68,10 @@ def add_table_storage_entry(
         "viewable": results_viewable,
         **metadata,
     }
+
+    # fix datatypes
+    entity["create_datetime"] = datetime.strptime(
+        str(entity["create_datetime"]), "%Y%m%d_%H%M%S"
+    ).replace(tzinfo=timezone.utc)
+
     table_client.create_entity(entity)

src/nhp/aci/status/model_run_status.py

@@ -28,19 +28,12 @@ def _get_progress_from_ats(
     try:
         entity = table_client.get_entity(partition_key=dataset, row_key=model_run_id)
 
-        res = {
+        return {
+            "status": entity.get("status", "unknown"),
             "complete": json.loads(entity.get("progress", "{}")),
             "model_runs": entity["model_runs"],
-            "container_group_name": None,
+            "container_group_name": entity.get("container_group_name", None),
         }
-        status = entity.get("status")
-
-        if status == "complete":
-            res["status"] = "complete"
-        else:
-            res["container_group_name"] = entity.get("container_group_name", None)
-
-        return res
     except ResourceNotFoundError:
         return {}
 
@@ -84,12 +77,14 @@ def get_model_run_status(
     if not status:
         return None
 
-    container_group_name = status.pop("container_group_name")
-    if not container_group_name:
+    if status["status"] == "complete":
         return status
 
+    if status["container_group_name"] is None:
+        return {**status, "state": "Creating"}
+
     try:
-        return {**status, **_get_aci_status(container_group_name, credential, config)}
+        return {**status, **_get_aci_status(status["container_group_name"], credential, config)}
 
     except ResourceNotFoundError:
         return {**status, "state": "Creating"}

tests/unit/nhp/aci/run_model/test_run_model_storage.py

@@ -1,5 +1,7 @@
 """Test the storage module in the nhp.aci.run_model package."""
 
+from datetime import datetime, timezone
+
 from azure.core.exceptions import ResourceExistsError
 
 from nhp.aci.run_model.storage import add_table_storage_entry, upload_params_to_blob
@@ -78,6 +80,7 @@ def test_add_table_storage_entry(mocker, config):
         "scenario": "scenario",
         "app_version": "v5.0",
         "model_run_id": "model-run-id",
+        "create_datetime": "20260518_120000",
     }
 
     # act
@@ -95,6 +98,8 @@ def test_add_table_storage_entry(mocker, config):
         "modelruns",
         credential="credential",
     )
+
+    metadata["create_datetime"] = datetime(2026, 5, 18, 12, 0, tzinfo=timezone.utc)  # ty: ignore[invalid-assignment]
     m_table_client().create_entity.assert_called_once_with(
         {
             "PartitionKey": "dataset",

tests/unit/nhp/aci/status/test_status_model_run_status.py

@@ -27,9 +27,10 @@ def test__get_progress_from_ats(mocker, config):
 
     # assert
     assert actual == {
+        "status": "running",
         "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
-        "model_runs": 100,
         "container_group_name": "name",
+        "model_runs": 100,
     }
 
     m_tc.assert_called_once_with(
@@ -108,6 +109,7 @@ def test_get_model_run_status(mocker, config):
     m_gpa = mocker.patch(
         "nhp.aci.status.model_run_status._get_progress_from_ats",
         return_value={
+            "status": "running",
             "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
             "model_runs": 100,
             "container_group_name": "name",
@@ -122,8 +124,10 @@ def test_get_model_run_status(mocker, config):
 
     # assert
     assert actual == {
+        "status": "running",
         "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
         "model_runs": 100,
+        "container_group_name": "name",
         "state": "running",
     }
     m_gpa.assert_called_once_with("dataset", "id", "credential", config)
@@ -145,6 +149,7 @@ def test_get_model_run_status_creates_credential_and_config_if_none(mocker, conf
     m_gpa = mocker.patch(
         "nhp.aci.status.model_run_status._get_progress_from_ats",
         return_value={
+            "status": "running",
             "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
             "model_runs": 100,
             "container_group_name": "name",
@@ -159,8 +164,10 @@ def test_get_model_run_status_creates_credential_and_config_if_none(mocker, conf
 
     # assert
     assert actual == {
+        "status": "running",
         "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
         "model_runs": 100,
+        "container_group_name": "name",
         "state": "running",
     }
     m_gpa.assert_called_once_with("dataset", "id", "credential", config)
@@ -175,6 +182,7 @@ def test_get_model_run_status_resource_not_found_with_status(mocker, config):
     mocker.patch(
         "nhp.aci.status.model_run_status._get_progress_from_ats",
         return_value={
+            "status": "running",
             "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
             "model_runs": 100,
             "container_group_name": "name",
@@ -195,8 +203,46 @@ def test_get_model_run_status_resource_not_found_with_status(mocker, config):
 
     # assert
     assert actual == {
+        "status": "running",
+        "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
+        "model_runs": 100,
+        "container_group_name": "name",
+        "state": "Creating",
+    }
+
+
+def test_get_model_run_status_returns_progress_without_aci_when_no_container_group_name(
+    mocker, config
+):
+    # arrange
+    mocker.patch(
+        "nhp.aci.status.model_run_status._get_progress_from_ats",
+        return_value={
+            "status": "running",
+            "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
+            "model_runs": 100,
+            "container_group_name": None,
+        },
+    )
+    mocker.patch(
+        "nhp.aci.status.model_run_status._get_aci_status",
+        side_effect=ResourceNotFoundError("Not Found"),
+    )
+
+    # act
+    actual = get_model_run_status(
+        "dataset",
+        "id",
+        "credential",  # ty: ignore[invalid-argument-type]
+        config,
+    )
+
+    # assert
+    assert actual == {
+        "status": "running",
         "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
         "model_runs": 100,
+        "container_group_name": None,
         "state": "Creating",
     }
 
@@ -224,9 +270,7 @@ def test_get_model_run_status_resource_not_found_without_status(mocker, config):
     assert actual is None
 
 
-def test_get_model_run_status_returns_progress_without_aci_when_no_container_group_name(
-    mocker, config
-):
+def test_get_model_run_status_returns_progress_without_aci_when_complete(mocker, config):
     # arrange
     m_gas = mocker.patch("nhp.aci.status.model_run_status._get_aci_status")
     mocker.patch(
@@ -251,6 +295,7 @@ def test_get_model_run_status_returns_progress_without_aci_when_no_container_gro
     assert actual == {
         "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
         "model_runs": 100,
+        "container_group_name": None,
         "status": "complete",
     }
     m_gas.assert_not_called()