Merge pull request #7448 from DIRACGridBot/cherry-pick-2-a048eb277-in…

…tegration

[sweep:integration] fix: AREX get pilot logging info with a token

src/DIRAC/Resources/Computing/AREXComputingElement.py

@@ -165,7 +165,7 @@ def _request(self, method, query, params=None, data=None, headers=None, timeout=
         except requests.RequestException as e:
             return S_ERROR(f"Request exception: {e}")
 
-    def _checkSession(self):
+    def _checkSession(self, mandatoryProxy: bool = False):
         """Check that the session exists and carries a valid proxy."""
         if not self.session:
             return S_ERROR("REST interface not initialised.")
@@ -174,23 +174,33 @@ def _checkSession(self):
         self.session.cert = None
         self.headers.pop("Authorization", None)
 
-        # Get a proxy: still mandatory, even if tokens are used to authenticate
-        if not self.proxy:
-            self.log.error("Proxy not set")
-            return S_ERROR("Proxy not set")
+        # We need a token or a proxy to interact with the REST interface
+        if not (self.token or self.proxy):
+            self.log.error("Proxy or token not set")
+            return S_ERROR("Proxy or token not set")
+
+        # A proxy might be required: in this case, it should be present
+        if mandatoryProxy and not self.proxy:
+            self.log.error("Proxy is mandatory but not set")
+            return S_ERROR("Proxy is mandatory but not set")
+
+        # If a proxy is required or if no token is set
+        if mandatoryProxy or not self.token:
+            # Prepare the proxy in X509_USER_PROXY
+            if not (result := self._prepareProxy())["OK"]:
+                self.log.error("Failed to set up proxy", result["Message"])
+                return result
 
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("Failed to set up proxy", result["Message"])
-            return result
+            # Attach the proxy to the session
+            self.session.cert = os.environ.get("X509_USER_PROXY")
+            self.log.verbose("A proxy is attached to the session")
 
+        # If a token is set, we use it
         if self.token:
             # Attach the token to the headers if present
-            self.headers["Authorization"] = "Bearer " + self.token["access_token"]
-            return S_OK()
+            self.headers["Authorization"] = f"Bearer {self.token['access_token']}"
+            self.log.verbose("A token is attached to the header of the request(s)")
 
-        # Attach the proxy to the session, only if the token is unavailable
-        self.session.cert = os.environ["X509_USER_PROXY"]
         return S_OK()
 
     #############################################################################
@@ -404,7 +414,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         Assume that the ARC queues are always of the format nordugrid-<batchSystem>-<queue>
         And none of our supported batch systems have a "-" in their name
         """
-        result = self._checkSession()
+        result = self._checkSession(mandatoryProxy=True)
         if not result["OK"]:
             self.log.error("Cannot submit jobs", result["Message"])
             return result
@@ -710,7 +720,7 @@ def getJobStatus(self, jobIDList):
 
         :param list jobIDList: list of job references, followed by the DIRAC stamp.
         """
-        result = self._checkSession()
+        result = self._checkSession(mandatoryProxy=True)
         if not result["OK"]:
             self.log.error("Cannot get status of the jobs", result["Message"])
             return result

src/DIRAC/Resources/Computing/test/Test_AREXComputingElement.py

@@ -0,0 +1,94 @@
+import os
+import requests
+from unittest.mock import MagicMock, patch
+
+from DIRAC import S_ERROR, S_OK
+
+
+# Assuming 'arc' is imported at the module level of ARCComputingElement
+@patch.dict("sys.modules", {"arc": MagicMock()})
+def test__checkSession(mocker):
+    """Test checkSession"""
+    from DIRAC.Resources.Computing.AREXComputingElement import AREXComputingElement
+
+    arex = AREXComputingElement("test")
+
+    # 1. The session has not been initialized: it shoud return an error
+    result = arex._checkSession()
+
+    assert not result["OK"], result
+    assert not arex.session
+    assert "Authorization" not in arex.headers, arex.headers
+
+    # 2. This time the session is initialized but there is no proxy nor token
+    # It should return an error
+    arex.session = requests.Session()
+    result = arex._checkSession()
+
+    assert not result["OK"], result
+    assert arex.session
+    assert not arex.session.cert
+    assert "Authorization" not in arex.headers, arex.headers
+
+    # 3. We set a malformed proxy, but not a token: it should return an error
+    arex.proxy = "fake proxy"
+    mocker.patch(
+        "DIRAC.Resources.Computing.AREXComputingElement.AREXComputingElement._prepareProxy", return_value=S_ERROR()
+    )
+
+    result = arex._checkSession()
+    assert not result["OK"], result
+    assert arex.session
+    assert not arex.session.cert
+    assert "Authorization" not in arex.headers, arex.headers
+
+    # 4. We set a proxy, but not a token: the session should include the proxy, but not the token
+    arex.proxy = "fake proxy"
+
+    def side_effect():
+        os.environ["X509_USER_PROXY"] = arex.proxy
+        return S_OK()
+
+    mocker.patch(
+        "DIRAC.Resources.Computing.AREXComputingElement.AREXComputingElement._prepareProxy", side_effect=side_effect
+    )
+
+    result = arex._checkSession()
+    assert result["OK"], result
+    assert arex.session
+    assert arex.session.cert
+    assert "Authorization" not in arex.headers, arex.headers
+
+    # 5. We set a proxy and a token: the session should just include
+    # the token because the proxy is not mandatory
+    arex.proxy = "fake proxy"
+    arex.token = {"access_token": "fake token"}
+
+    result = arex._checkSession()
+    assert result["OK"], result
+    assert arex.session
+    assert not arex.session.cert
+    assert "Authorization" in arex.headers, arex.headers
+
+    # 6. We make the proxy mandatory: the session should include both the proxy and the token
+    result = arex._checkSession(mandatoryProxy=True)
+    assert result["OK"], result
+    assert arex.session
+    assert arex.session.cert
+    assert "Authorization" in arex.headers, arex.headers
+
+    # 7. Now we just include the token, the proxy is not mandatory:
+    # the session should only include the token
+    arex.proxy = None
+    result = arex._checkSession()
+    assert result["OK"], result
+    assert arex.session
+    assert not arex.session.cert
+    assert "Authorization" in arex.headers, arex.headers
+
+    # 8. Now we just include the token, but the proxy is mandatory: it should return an error
+    result = arex._checkSession(mandatoryProxy=True)
+    assert not result["OK"], result
+    assert arex.session
+    assert not arex.session.cert
+    assert "Authorization" not in arex.headers, arex.headers