test: Adding tests to the WMS access policy, and some fixes

Author: Robin-Van-de-Merghel

diracx-routers/src/diracx/routers/jobs/access_policies.py

@@ -59,10 +59,6 @@ async def policy(
                 pilot_db
             ), "pilot_db is a mandatory parameter when using a pilot action"
             assert job_ids, "job_ids has to be defined"
-            assert (
-                len(job_ids) == 1
-            ), "a pilot can have only one job_id associated, and it has to be given"
-
             pilot_info = user_info  # For semantic
 
             # Syntax to avoid code duplication
@@ -83,11 +79,11 @@ async def policy(
                 )
 
                 # Equivalent of issubset, but cleaner
-                if set(job_ids) <= pilot_jobs:
+                if set(job_ids) <= set(pilot_jobs):
                     return
 
                 raise HTTPException(
-                    status.HTTP_403_FORBIDDEN, "this pilot can't modify this job"
+                    status.HTTP_403_FORBIDDEN, "this pilot can't access/modify this job"
                 )
 
             raise HTTPException(status.HTTP_403_FORBIDDEN, "you are not a pilot")

diracx-routers/tests/jobs/test_wms_access_policy.py

@@ -4,7 +4,7 @@
 from fastapi import HTTPException, status
 from uuid_utils import uuid7
 
-from diracx.core.properties import JOB_ADMINISTRATOR, NORMAL_USER
+from diracx.core.properties import GENERIC_PILOT, JOB_ADMINISTRATOR, NORMAL_USER
 from diracx.routers.jobs.access_policies import (
     ActionType,
     SandboxAccessPolicy,
@@ -26,6 +26,11 @@ class FakeJobDB:
     async def summary(self, *args): ...
 
 
+class FakePilotDB:
+    async def get_pilot_by_reference(self, *args): ...
+    async def get_pilot_job_ids(self, *args): ...
+
+
 class FakeSBMetadataDB:
     async def get_owner_id(self, *args): ...
     async def get_sandbox_owner_id(self, *args): ...
@@ -36,6 +41,11 @@ def job_db():
     yield FakeJobDB()
 
 
+@pytest.fixture
+def pilot_db():
+    yield FakePilotDB()
+
+
 @pytest.fixture
 def sandbox_metadata_db():
     yield FakeSBMetadataDB()
@@ -68,6 +78,112 @@ async def test_wms_access_policy_weird_user(job_db):
         )
 
 
+async def test_wms_access_policy_pilot(job_db, pilot_db, monkeypatch):
+
+    normal_user = AuthorizedUserInfo(properties=[NORMAL_USER], **base_payload)
+    pilot = AuthorizedUserInfo(properties=[GENERIC_PILOT], **base_payload)
+
+    # ------------------------- Simple User accessing a pilot action -------------------------
+    # A user cannot create any resource
+    with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
+        await WMSAccessPolicy.policy(
+            WMS_POLICY_NAME,
+            normal_user,
+            action=ActionType.PILOT,
+            job_db=job_db,
+            pilot_db=pilot_db,
+            job_ids=[1, 2],
+        )
+
+    # Split to distinguish the generated part ("403 ") from the message part ("you are not a pilot")
+    assert str(excinfo.value) == "403: " + "you are not a pilot", excinfo
+
+    # ------------------------- Lost pilot -------------------------
+    async def get_pilot_by_reference_patch(*args):
+        return []
+
+    monkeypatch.setattr(
+        pilot_db, "get_pilot_by_reference", get_pilot_by_reference_patch
+    )
+
+    # A pilot that has expired (removed from db) should not be able to access jobs
+    with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
+        await WMSAccessPolicy.policy(
+            WMS_POLICY_NAME,
+            pilot,
+            action=ActionType.PILOT,
+            pilot_db=pilot_db,
+            job_db=job_db,
+            job_ids=[1, 2],
+        )
+
+    assert str(excinfo.value) == "403: " + "this pilot is not registered", excinfo
+
+    # ------------------------- Pilot accessing wrong jobs -------------------------
+    async def get_pilot_by_reference_patch(*args, **kwargs):
+        return {"PilotID": 1}
+
+    async def get_pilot_job_ids_patch(*args, **kwargs):
+        return []
+
+    monkeypatch.setattr(
+        pilot_db, "get_pilot_by_reference", get_pilot_by_reference_patch
+    )
+    monkeypatch.setattr(pilot_db, "get_pilot_job_ids", get_pilot_job_ids_patch)
+
+    # A pilot that has is not associated with a job can't access a job
+    with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
+        await WMSAccessPolicy.policy(
+            WMS_POLICY_NAME,
+            pilot,
+            action=ActionType.PILOT,
+            pilot_db=pilot_db,
+            job_db=job_db,
+            job_ids=[1, 2],
+        )
+
+    assert (
+        str(excinfo.value) == "403: " + "this pilot can't access/modify this job"
+    ), excinfo
+
+    # ------------------------- Pilot accessing some of his jobs -------------------------
+    async def get_pilot_job_ids_patch(*args, **kwargs):
+        return [1, 2, 3, 4]
+
+    monkeypatch.setattr(pilot_db, "get_pilot_job_ids", get_pilot_job_ids_patch)
+
+    # A pilot that is associated with a job can access a job
+    await WMSAccessPolicy.policy(
+        WMS_POLICY_NAME,
+        pilot,
+        action=ActionType.PILOT,
+        pilot_db=pilot_db,
+        job_db=job_db,
+        job_ids=[1, 2],
+    )
+
+    # ------------------------- Pilot accessing some of his jobs plus some forbidden -------------------------
+    async def get_pilot_job_ids_patch(*args, **kwargs):
+        return [1, 2, 3, 4]
+
+    monkeypatch.setattr(pilot_db, "get_pilot_job_ids", get_pilot_job_ids_patch)
+
+    # A pilot that fetches few jobs, one where he does not have the rights, and few where he has the rights
+    with pytest.raises(HTTPException, match=f"{status.HTTP_403_FORBIDDEN}") as excinfo:
+        await WMSAccessPolicy.policy(
+            WMS_POLICY_NAME,
+            pilot,
+            action=ActionType.PILOT,
+            pilot_db=pilot_db,
+            job_db=job_db,
+            job_ids=[1, 2, 12],
+        )
+
+    assert (
+        str(excinfo.value) == "403: " + "this pilot can't access/modify this job"
+    ), excinfo
+
+
 async def test_wms_access_policy_create(job_db):
 
     admin_user = AuthorizedUserInfo(properties=[JOB_ADMINISTRATOR], **base_payload)

diracx-testing/src/diracx/testing/utils.py

@@ -359,6 +359,29 @@ def unauthenticated(self):
         with TestClient(self.app) as client:
             yield client
 
+    @contextlib.contextmanager
+    def pilot(self):
+        from diracx.core.properties import GENERIC_PILOT, LIMITED_DELEGATION
+        from diracx.routers.auth.token import create_token
+
+        with self.unauthenticated() as client:
+            payload = {
+                "sub": "testingVO:yellow-sub",
+                "exp": datetime.now(tz=timezone.utc)
+                + timedelta(self.test_auth_settings.access_token_expire_minutes),
+                "iss": ISSUER,
+                "dirac_properties": [GENERIC_PILOT, LIMITED_DELEGATION],
+                "jti": str(uuid4()),
+                "preferred_username": "preferred_username",
+                "dirac_group": "test_group",
+                "vo": "lhcb",
+            }
+            token = create_token(payload, self.test_auth_settings)
+
+            client.headers["Authorization"] = f"Bearer {token}"
+            client.dirac_token_payload = payload
+            yield client
+
     @contextlib.contextmanager
     def normal_user(self):
         from diracx.core.properties import NORMAL_USER