Add account parameter in job_info method (#196)

ekouts · web-flow · commit c03c2c485697 · 2025-12-09T11:30:53.000+01:00
diff --git a/firecrest/cli2/__init__.py b/firecrest/cli2/__init__.py
@@ -854,11 +854,19 @@ def job_info(
     jobid: Optional[str] = typer.Argument(
          None, help="Job ID to display."
     ),
+    allusers: bool = typer.Option(
+        False, "--all-users", help="Retrieve information about all users' jobs."
+    ),
+    account: Optional[str] = typer.Option(
+        None,
+        help=("An account to filter the jobs by. It will only be taken into "
+              "account when you are not specifying a jobid."),
+    ),
 ):
     """Retrieve information about submitted jobs.
     """
     try:
-        result = asyncio.run(client.job_info(system, jobid))
+        result = asyncio.run(client.job_info(system, jobid, allusers, account))
         json_out(result)
     except Exception as e:
         examine_exeption(e)
diff --git a/firecrest/v2/_async/Client.py b/firecrest/v2/_async/Client.py
@@ -1646,7 +1646,8 @@ async def job_info(
         self,
         system_name: str,
         jobid: Optional[str] = None,
-        allusers: bool = False
+        allusers: bool = False,
+        account: Optional[str] = None
     ) -> list:
         """Get job information. When the job is not specified, it will return
         all the jobs.
@@ -1655,6 +1656,8 @@ async def job_info(
         :param jobid: the ID of the job
         :param allusers: whether to return jobs of all users or only the
                          current user
+        :param account: an account to filter the jobs by. It will only be taken
+                        into account when you are not specifying a jobid.
         :calls: GET `/compute/{system_name}/jobs` or
                 GET `/compute/{system_name}/jobs/{job}`
         """
@@ -1667,9 +1670,15 @@ async def job_info(
                 "version <2.2.7 of the API."
             )
 
+        if self._api_version < parse("2.4.2") and account:
+            raise NotImplementedOnAPIversion(
+                "The `account` parameter is not available for "
+                "version <2.4.2 of the API."
+            )
+
         resp = await self._get_request(
             endpoint=url,
-            params={"allusers": allusers}
+            params={"allusers": allusers, "account": account}
         )
         result_jobs = self._check_response(resp, 200)["jobs"]
         return result_jobs if result_jobs is not None else []
diff --git a/firecrest/v2/_sync/Client.py b/firecrest/v2/_sync/Client.py
@@ -1645,7 +1645,8 @@ def job_info(
         self,
         system_name: str,
         jobid: Optional[str] = None,
-        allusers: bool = False
+        allusers: bool = False,
+        account: Optional[str] = None
     ) -> list:
         """Get job information. When the job is not specified, it will return
         all the jobs.
@@ -1654,6 +1655,8 @@ def job_info(
         :param jobid: the ID of the job
         :param allusers: whether to return jobs of all users or only the
                          current user
+        :param account: an account to filter the jobs by. It will only be taken
+                        into account when you are not specifying a jobid.
         :calls: GET `/compute/{system_name}/jobs` or
                 GET `/compute/{system_name}/jobs/{job}`
         """
@@ -1666,9 +1669,15 @@ def job_info(
                 "version <2.2.7 of the API."
             )
 
+        if self._api_version < parse("2.4.2") and account:
+            raise NotImplementedOnAPIversion(
+                "The `account` parameter is not available for "
+                "version <2.4.2 of the API."
+            )
+
         resp = self._get_request(
             endpoint=url,
-            params={"allusers": allusers}
+            params={"allusers": allusers, "account": account}
         )
         result_jobs = self._check_response(resp, 200)["jobs"]
         return result_jobs if result_jobs is not None else []
diff --git a/tests/v2/handlers.py b/tests/v2/handlers.py
@@ -47,7 +47,7 @@ def fc_server(httpserver):
     for endpoint in ["jobs", "metadata"]:
         httpserver.expect_request(
             re.compile(rf"/compute/.*/{endpoint}"), method="GET"
-        ).respond_with_handler(filesystem_handler)
+        ).respond_with_handler(get_jobs_handler)
 
     httpserver.expect_request(
         re.compile(r"/compute/.*/jobs"), method="POST"
@@ -268,3 +268,45 @@ def submit_handler(request: Request):
     return Response(json.dumps(ret),
                     status=ret_status,
                     content_type="application/json")
+
+
+def get_jobs_handler(request: Request):
+    if request.headers["Authorization"] != "Bearer VALID_TOKEN":
+        return Response(
+            json.dumps({"message": "Bad token; invalid JSON"}),
+            status=401,
+            content_type="application/json",
+        )
+
+    url, *params = request.url.split("?")
+
+    endpoint = url.split("/")[-1]
+
+    suffix = ""
+
+    if endpoint == "jobs":
+        endpoint = "job"
+
+        if "account=users2" in "&".join(params):
+            suffix = "_info_account"
+        elif "allusers=true" in "&".join(params):
+            suffix = "_info_all_users"
+        else:
+            suffix = "_info"
+
+    if endpoint == "1":
+        endpoint = "job"
+        suffix = "_info"
+
+    if endpoint == "metadata":
+        endpoint = "job"
+        suffix = "_metadata"
+
+    data = read_json_file(f"v2/responses/{endpoint}{suffix}.json")
+
+    ret = data["response"]
+    ret_status = data["status_code"]
+
+    return Response(json.dumps(ret),
+                    status=ret_status,
+                    content_type="application/json")
diff --git a/tests/v2/responses/job_info_account.json b/tests/v2/responses/job_info_account.json
@@ -0,0 +1,41 @@
+{
+    "status_code": 200,
+    "response": {
+        "jobs": [
+            {
+                "account": "test",
+                "allocationNodes": 1,
+                "cluster": "cluster",
+                "exitCode": {
+                    "status": [
+                        "ERROR"
+                    ],
+                    "returnCode": 8
+                },
+                "group": "users2",
+                "jobId": 3,
+                "name": "allocation",
+                "nodes": "localhost",
+                "partition": "part01",
+                "priority": 1,
+                "killRequestUser": "",
+                "state": {
+                    "current": [
+                        "FAILED"
+                    ],
+                    "reason": "None"
+                },
+                "time": {
+                    "elapsed": 0,
+                    "start": 1730968675,
+                    "end": 1730968675,
+                    "suspended": 0,
+                    "limit": 7200,
+                    "submission": 1730968674
+                },
+                "user": "test3",
+                "workingDirectory": "/home/test3"
+            }
+        ]
+    }
+}
diff --git a/tests/v2/responses/job_info_all_users.json b/tests/v2/responses/job_info_all_users.json
@@ -0,0 +1,109 @@
+{
+    "status_code": 200,
+    "response": {
+        "jobs": [
+            {
+                "account": "test",
+                "allocationNodes": 1,
+                "cluster": "cluster",
+                "exitCode": {
+                    "status": [
+                        "ERROR"
+                    ],
+                    "returnCode": 8
+                },
+                "group": "users",
+                "jobId": 1,
+                "name": "allocation",
+                "nodes": "localhost",
+                "partition": "part01",
+                "priority": 1,
+                "killRequestUser": "",
+                "state": {
+                    "current": [
+                        "FAILED"
+                    ],
+                    "reason": "None"
+                },
+                "time": {
+                    "elapsed": 0,
+                    "start": 1730968675,
+                    "end": 1730968675,
+                    "suspended": 0,
+                    "limit": 7200,
+                    "submission": 1730968674
+                },
+                "user": "test1",
+                "workingDirectory": "/home/test1"
+            },
+            {
+                "account": "test",
+                "allocationNodes": 2,
+                "cluster": "cluster",
+                "exitCode": {
+                    "status": [
+                        "ERROR"
+                    ],
+                    "returnCode": 8
+                },
+                "group": "users",
+                "jobId": 2,
+                "name": "allocation",
+                "nodes": "localhost",
+                "partition": "part01",
+                "priority": 1,
+                "killRequestUser": "",
+                "state": {
+                    "current": [
+                        "FAILED"
+                    ],
+                    "reason": "None"
+                },
+                "time": {
+                    "elapsed": 0,
+                    "start": 1730968675,
+                    "end": 1730968675,
+                    "suspended": 0,
+                    "limit": 7200,
+                    "submission": 1730968674
+                },
+                "user": "test2",
+                "workingDirectory": "/home/test2"
+            },
+            {
+                "account": "test",
+                "allocationNodes": 1,
+                "cluster": "cluster",
+                "exitCode": {
+                    "status": [
+                        "ERROR"
+                    ],
+                    "returnCode": 8
+                },
+                "group": "users2",
+                "jobId": 3,
+                "name": "allocation",
+                "nodes": "localhost",
+                "partition": "part01",
+                "priority": 1,
+                "killRequestUser": "",
+                "state": {
+                    "current": [
+                        "FAILED"
+                    ],
+                    "reason": "None"
+                },
+                "time": {
+                    "elapsed": 0,
+                    "start": 1730968675,
+                    "end": 1730968675,
+                    "suspended": 0,
+                    "limit": 7200,
+                    "submission": 1730968674
+                },
+                "user": "test3",
+                "workingDirectory": "/home/test3"
+            }
+        ]
+    }
+}
diff --git a/tests/v2/test_v2_async.py b/tests/v2/test_v2_async.py
@@ -1,6 +1,5 @@
 import json
 import pytest
-import re
 
 from context_v2 import AsyncFirecrest, UnexpectedStatusException
 from werkzeug.wrappers import Response
@@ -18,10 +17,12 @@ class ValidAuthorization:
         def get_access_token(self):
             return "VALID_TOKEN"
 
-    return AsyncFirecrest(
+    client = AsyncFirecrest(
         firecrest_url=fc_server.url_for("/"),
         authorization=ValidAuthorization()
     )
+    client.set_api_version("2.99.0") # Set to a high enough version for all tests
+    return client
 
 
 @pytest.fixture
@@ -317,6 +318,20 @@ async def test_job_info_jobid(valid_client):
     assert resp == data["response"]["jobs"]
 
 
+@pytest.mark.asyncio
+async def test_job_allusers(valid_client):
+    data = read_json_file("v2/responses/job_info_all_users.json")
+    resp = await valid_client.job_info("cluster", allusers=True)
+    assert resp == data["response"]["jobs"]
+
+
+@pytest.mark.asyncio
+async def test_job_info_account(valid_client):
+    data = read_json_file("v2/responses/job_info_account.json")
+    resp = await valid_client.job_info("cluster", account="users2")
+    assert resp == data["response"]["jobs"]
+
+
 @pytest.mark.asyncio
 async def test_job_metadata(valid_client):
     data = read_json_file("v2/responses/job_metadata.json")
diff --git a/tests/v2/test_v2_sync.py b/tests/v2/test_v2_sync.py
@@ -1,6 +1,5 @@
 import json
 import pytest
-import re
 
 from context_v2 import Firecrest, UnexpectedStatusException
 from werkzeug.wrappers import Response
@@ -18,11 +17,12 @@ class ValidAuthorization:
         def get_access_token(self):
             return "VALID_TOKEN"
 
-    return Firecrest(
+    client = Firecrest(
         firecrest_url=fc_server.url_for("/"),
         authorization=ValidAuthorization()
     )
-
+    client.set_api_version("2.99.0") # Set to a high enough version for all tests
+    return client
 
 @pytest.fixture
 def invalid_client(fc_server):
@@ -293,6 +293,20 @@ def test_job_info_jobid(valid_client):
     assert resp == data["response"]["jobs"]
 
 
+def test_job_allusers(valid_client):
+    data = read_json_file("v2/responses/job_info_all_users.json")
+    resp = valid_client.job_info("cluster", allusers=True)
+
+    assert resp == data["response"]["jobs"]
+
+
+def test_job_info_account(valid_client):
+    data = read_json_file("v2/responses/job_info_account.json")
+    resp = valid_client.job_info("cluster", account="users2")
+
+    assert resp == data["response"]["jobs"]
+
+
 def test_job_metadata(valid_client):
     data = read_json_file("v2/responses/job_metadata.json")
     resp = valid_client.job_metadata("cluster", "1")
