cherry-pick: Pull commits from release 5.5 (#726)

Author: dusktreader

changes/agent/.gitignore

@@ -0,0 +1 @@
+Updated the schema for slurm job data to handle slurm updates

changes/agent/722.fixed.md

@@ -2,7 +2,7 @@
 from typing import List, Optional, TypedDict, TypeAlias
 
 import pydantic
-from pydantic import ConfigDict
+from pydantic import ConfigDict, field_validator
 
 from jobbergate_agent.jobbergate.constants import FileType, INFLUXDB_MEASUREMENT
 
@@ -76,9 +76,24 @@ class SlurmJobData(pydantic.BaseModel, extra="ignore"):
 
     job_id: Optional[int] = None
     job_state: Optional[str] = None
-    job_info: Optional[str] = None
+    job_info: Optional[str] = "{}"
     state_reason: Optional[str] = None
 
+    @field_validator("job_state", mode="before")
+    @classmethod
+    def validate_job_state(cls, value: str | list[str] | None) -> str | None:
+        """
+        Validate the job_state field.
+        """
+        if value is None:
+            return None
+
+        if isinstance(value, list):
+            if len(value) != 1:
+                raise ValueError("job_state does not have exactly one value.")
+            return value[0]
+        return value
+
 
 class InfluxDBMeasurementDict(TypedDict):
     """

jobbergate-agent/jobbergate_agent/jobbergate/schemas.py

@@ -24,7 +24,7 @@
 )
 from jobbergate_agent.jobbergate.constants import INFLUXDB_MEASUREMENT
 from jobbergate_agent.settings import SETTINGS
-from jobbergate_agent.utils.exception import JobbergateApiError, JobbergateAgentError
+from jobbergate_agent.utils.exception import JobbergateApiError, JobbergateAgentError, SbatchError
 
 
 @pytest.fixture()
@@ -84,6 +84,55 @@ async def test_fetch_job_data__success():
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_access_token")
+async def test_fetch_job_data__handles_list_in_job_state():
+    """
+    Test that the ``fetch_job_data()`` function can successfully retrieve
+    job data from Slurm as a ``SlurmJobData`` when the job_state from slurm
+    is reported as a list.
+    """
+    mocked_sbatch = mock.MagicMock()
+    mocked_sbatch.get_job_info.return_value = dict(
+        job_state=["FAILED"],
+        job_id=123,
+        state_reason="NonZeroExitCode",
+        foo="bar",
+    )
+
+    result: SlurmJobData = await fetch_job_data(123, mocked_sbatch)
+
+    assert result.job_id == 123
+    assert result.job_state == "FAILED"
+    assert result.state_reason == "NonZeroExitCode"
+    assert result.job_info is not None
+    assert json.loads(result.job_info) == dict(
+        job_state=["FAILED"],
+        job_id=123,
+        state_reason="NonZeroExitCode",
+        foo="bar",
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_access_token")
+async def test_fetch_job_data__raises_error_if_job_state_is_invalid_list():
+    """
+    Test that the ``fetch_job_data()`` function raises an exception
+    if the list slurm_job_state does not have exactly one value.
+    """
+    mocked_sbatch = mock.MagicMock()
+    mocked_sbatch.get_job_info.return_value = dict(
+        job_state=[],
+        job_id=123,
+        state_reason="NonZeroExitCode",
+        foo="bar",
+    )
+
+    with pytest.raises(SbatchError, match="does not have exactly one value"):
+        await fetch_job_data(123, mocked_sbatch)
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("mock_access_token")
 async def test_fetch_job_data__reports_status_as_UNKOWN_if_slurm_job_id_is_not_found():