improves the logic for getting the status of completed model runs

src/nhp/aci/__main__.py

@@ -97,6 +97,14 @@ def _status_single_model_run(dataset: str, model_run_id: str) -> None:
         print(f"Unknown model run id: {model_run_id}")
         return
 
+    if "status" in status and status["status"] == "complete":
+        print(f"{model_run_id}: complete")
+        return
+
+    if "state" not in status:
+        print(f"Model run {model_run_id} has no state information available")
+        return
+
     state = status["state"]
 
     complete = status["complete"]

src/nhp/aci/status/model_run_status.py

@@ -28,11 +28,19 @@ def _get_progress_from_ats(
     try:
         entity = table_client.get_entity(partition_key=dataset, row_key=model_run_id)
 
-        return {
-            "complete": json.loads(entity["progress"]),
+        res = {
+            "complete": json.loads(entity.get("progress", "{}")),
             "model_runs": entity["model_runs"],
-            "container_group_name": entity["container_group_name"],
+            "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 {}
 
@@ -77,6 +85,8 @@ def get_model_run_status(
         return None
 
     container_group_name = status.pop("container_group_name")
+    if not container_group_name:
+        return status
 
     try:
         return {**status, **_get_aci_status(container_group_name, credential, config)}

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

@@ -16,6 +16,7 @@ def test__get_progress_from_ats(mocker, config):
     m_tc.return_value.get_entity.return_value = {
         "progress": '{"inpatients": 100, "outpatients": 50, "aae": 0}',
         "model_runs": 100,
+        "status": "running",
         "container_group_name": "name",
     }
 
@@ -57,6 +58,28 @@ def test__get_progress_from_ats_not_found(mocker, config):
     m_tc.return_value.get_entity.assert_called_once_with(partition_key="name", row_key="id")
 
 
+def test__get_progress_from_ats_complete_status_without_container_group_name(mocker, config):
+    # arrange
+    m_tc = mocker.patch("nhp.aci.status.model_run_status.TableClient")
+
+    m_tc.return_value.get_entity.return_value = {
+        "model_runs": 100,
+        "status": "complete",
+    }
+
+    # act
+    actual = _get_progress_from_ats("dataset", "id", "credential", config)  # ty: ignore[invalid-argument-type]
+
+    # assert
+    assert actual == {
+        "complete": {},
+        "model_runs": 100,
+        "container_group_name": None,
+        "status": "complete",
+    }
+    m_tc.return_value.get_entity.assert_called_once_with(partition_key="dataset", row_key="id")
+
+
 def test__get_aci_status(mocker, config):
     # arrange
     m_cimc = mocker.patch("nhp.aci.status.model_run_status.ContainerInstanceManagementClient")
@@ -199,3 +222,35 @@ def test_get_model_run_status_resource_not_found_without_status(mocker, config):
 
     # assert
     assert actual is None
+
+
+def test_get_model_run_status_returns_progress_without_aci_when_no_container_group_name(
+    mocker, config
+):
+    # arrange
+    m_gas = mocker.patch("nhp.aci.status.model_run_status._get_aci_status")
+    mocker.patch(
+        "nhp.aci.status.model_run_status._get_progress_from_ats",
+        return_value={
+            "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
+            "model_runs": 100,
+            "container_group_name": None,
+            "status": "complete",
+        },
+    )
+
+    # act
+    actual = get_model_run_status(
+        "dataset",
+        "id",
+        "credential",  # ty: ignore[invalid-argument-type]
+        config,
+    )
+
+    # assert
+    assert actual == {
+        "complete": {"inpatients": 100, "outpatients": 50, "aae": 0},
+        "model_runs": 100,
+        "status": "complete",
+    }
+    m_gas.assert_not_called()

tests/unit/nhp/aci/test___main__.py

@@ -285,6 +285,38 @@ def test_single_model_run_not_exists(self, mocker):
         m_status.assert_called_once_with("dataset", "id")
         m_print.assert_called_once_with("Unknown model run id: id")
 
+    def test_single_model_run_complete_status(self, mocker):
+        # arrange
+        status = {"status": "complete"}
+        m_status = mocker.patch("nhp.aci.__main__.get_model_run_status", return_value=status)
+        args = Mock()
+        args.dataset = "dataset"
+        args.model_id = "id"
+        m_print = mocker.patch("builtins.print")
+
+        # act
+        _status(args)
+
+        # assert
+        m_status.assert_called_once_with("dataset", "id")
+        m_print.assert_called_once_with("id: complete")
+
+    def test_single_model_run_missing_state(self, mocker):
+        # arrange
+        status = {"status": "running"}
+        m_status = mocker.patch("nhp.aci.__main__.get_model_run_status", return_value=status)
+        args = Mock()
+        args.dataset = "dataset"
+        args.model_id = "id"
+        m_print = mocker.patch("builtins.print")
+
+        # act
+        _status(args)
+
+        # assert
+        m_status.assert_called_once_with("dataset", "id")
+        m_print.assert_called_once_with("Model run id has no state information available")
+
     def test_all_model_runs_no_runs(self, mocker):
         # arrange
         m_cmr = mocker.patch("nhp.aci.__main__.get_current_model_runs", return_value={})