From 8cf7cda9ce1be00d514b4912a1d9d5648aa32c29 Mon Sep 17 00:00:00 2001
From: Federico Stagni <federico.stagni@cern.ch>
Date: Mon, 18 Dec 2023 15:13:37 +0100
Subject: [PATCH] feat: using CVMFS_locations for discovering the pilot files

---
 .../Utilities/PilotWrapper.py                 | 38 ++++++++++++++-----
 .../Utilities/test/Test_PilotWrapper.py       | 36 +++++++++++++++++-
 2 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py b/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py
index a2fe5d0a372..192affc9905 100644
--- a/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py
+++ b/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py
@@ -11,14 +11,12 @@
         _writePilotWrapperFile(localPilot=localPilot)
 
 """
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
-import os
-import tempfile
 import base64
 import bz2
+import os
+import tempfile
 
 pilotWrapperContent = """#!/bin/bash
 if command -v python &> /dev/null; then
@@ -78,7 +76,12 @@
 
 
 def pilotWrapperScript(
-    pilotFilesCompressedEncodedDict=None, pilotOptions="", pilotExecDir="", envVariables=None, location=""
+    pilotFilesCompressedEncodedDict=None,
+    pilotOptions="",
+    pilotExecDir="",
+    envVariables=None,
+    location="",
+    CVMFS_locations=None,
 ):
     """Returns the content of the pilot wrapper script.
 
@@ -95,6 +98,8 @@ def pilotWrapperScript(
     :type envVariables: dict
     :param location: location where to get the pilot files
     :type location: string
+    :param location: optional CVMFS locations of where to get the pilot files
+    :type location: list
 
     :returns: content of the pilot wrapper
     :rtype: string
@@ -106,6 +111,20 @@ def pilotWrapperScript(
     if envVariables is None:
         envVariables = {}
 
+    if CVMFS_locations is None:
+        # What is in this location is almost certainly incorrect, especially the pilot.json
+        CVMFS_locs = '["file:/cvmfs/dirac.egi.eu/pilot"]'
+    else:
+        # Here we are making the assumption that, if CVMFS_locations is, e.g., ['/cvmfs/somewhere', '/cvmfs/elsewhere']
+        # and the project is 'LHCb',
+        # then the pilot can maybe be found at locations
+        # - file:/cvmfs/somewhere/lhcbdirac/pilot
+        # - file:/cvmfs/elsewhere/lhcbdirac/pilot
+        project = "dirac"
+        if "-l" in pilotOptions:
+            project = pilotOptions.split(" ")[pilotOptions.split(" ").index("-l") + 1].lower() + "dirac"
+        CVMFS_locs = "[" + ",".join('"file:' + os.path.join(loc, project, 'pilot"') for loc in CVMFS_locations) + "]"
+
     compressedString = ""
     # are there some pilot files to unpack? Then we create the unpacking string
     for pfName, encodedPf in pilotFilesCompressedEncodedDict.items():
@@ -178,8 +197,8 @@ def pilotWrapperScript(
 # we try from the available locations
 locs = [os.path.join('https://', loc) for loc in location]
 locations = locs + [os.path.join(loc, 'pilot') for loc in locs]
-# adding also, as last the cvmfs location dirac.egi.eu, but this won't contain a valid JSON
-locations += ['file:/cvmfs/dirac.egi.eu/pilot/']
+# adding also the cvmfs locations
+locations += %(CVMFS_locs)s
 
 for loc in locations:
   print('Trying %%s' %% loc)
@@ -262,7 +281,8 @@ def pilotWrapperScript(
     logger.debug('Checksum matched')
 
 """ % {
-            "location": location
+            "location": location,
+            "CVMFS_locs": CVMFS_locs,
         }
 
     localPilot += (
diff --git a/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py b/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py
index 0b7f33f3c3d..96cec570f8b 100644
--- a/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py
+++ b/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py
@@ -1,9 +1,9 @@
 """ This is a test of the creation of the pilot wrapper
 """
 
-import os
 import base64
 import bz2
+import os
 
 from DIRAC.WorkloadManagementSystem.Utilities.PilotWrapper import pilotWrapperScript
 
@@ -87,3 +87,37 @@ def test_scriptPilot3():
 
     assert 'os.environ["someName"]="someValue"' in res
     assert "lhcb-portal.cern.ch" in res
+    assert """locations += ["file:/cvmfs/dirac.egi.eu/pilot"]""" in res
+
+
+def test_scriptPilot3_2():
+    """test script creation"""
+    res = pilotWrapperScript(
+        pilotFilesCompressedEncodedDict={"proxy": "thisIsSomeProxy"},
+        pilotOptions="-c 123 --foo bar",
+        envVariables={"someName": "someValue", "someMore": "oneMore"},
+        location="lhcb-portal.cern.ch",
+        CVMFS_locations=["/cvmfs/lhcb.cern.ch", "/cvmfs/dirac.egi.eu"],
+    )
+
+    assert 'os.environ["someName"]="someValue"' in res
+    assert "lhcb-portal.cern.ch" in res
+    assert """locations += ["file:/cvmfs/lhcb.cern.ch/dirac/pilot","file:/cvmfs/dirac.egi.eu/dirac/pilot"]""" in res
+
+
+def test_scriptPilot3_3():
+    """test script creation"""
+    res = pilotWrapperScript(
+        pilotFilesCompressedEncodedDict={"proxy": "thisIsSomeProxy"},
+        pilotOptions="-c 123 --foo bar -l LHCb -h pippo",
+        envVariables={"someName": "someValue", "someMore": "oneMore"},
+        location="lhcb-portal.cern.ch",
+        CVMFS_locations=["/cvmfs/lhcb.cern.ch", "/cvmfs/dirac.egi.eu"],
+    )
+
+    assert 'os.environ["someName"]="someValue"' in res
+    assert "lhcb-portal.cern.ch" in res
+    assert (
+        """locations += ["file:/cvmfs/lhcb.cern.ch/lhcbdirac/pilot","file:/cvmfs/dirac.egi.eu/lhcbdirac/pilot"]"""
+        in res
+    )