From ca7a2109164669cd0bbfb6ad234d9203d56324e6 Mon Sep 17 00:00:00 2001 From: aldbr Date: Wed, 7 Feb 2024 10:43:49 +0100 Subject: [PATCH] fix: AREXCE getLoggingInfo with a token --- .../Computing/AREXComputingElement.py | 40 +++++--- .../test/Test_AREXComputingElement.py | 94 +++++++++++++++++++ 2 files changed, 119 insertions(+), 15 deletions(-) create mode 100644 src/DIRAC/Resources/Computing/test/Test_AREXComputingElement.py diff --git a/src/DIRAC/Resources/Computing/AREXComputingElement.py b/src/DIRAC/Resources/Computing/AREXComputingElement.py index bf8e395db28..18815da59f6 100755 --- a/src/DIRAC/Resources/Computing/AREXComputingElement.py +++ b/src/DIRAC/Resources/Computing/AREXComputingElement.py @@ -174,7 +174,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.") @@ -183,23 +183,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() ############################################################################# @@ -413,7 +423,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs= Assume that the ARC queues are always of the format nordugrid-- 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 @@ -715,7 +725,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 diff --git a/src/DIRAC/Resources/Computing/test/Test_AREXComputingElement.py b/src/DIRAC/Resources/Computing/test/Test_AREXComputingElement.py new file mode 100644 index 00000000000..eeb22e21d3d --- /dev/null +++ b/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