From 7b5c00fd0a418520a90549b74781cf8584de75fc Mon Sep 17 00:00:00 2001
From: Andrea Piccinelli <andrea8395@gmail.com>
Date: Fri, 13 Dec 2024 16:57:35 +0100
Subject: [PATCH] Fixing unit tests

---
 .../Storage_t/Backends_t/GFAL2Impl_t.py       | 65 ++++++++++++++-----
 1 file changed, 48 insertions(+), 17 deletions(-)

diff --git a/test/python/WMCore_t/Storage_t/Backends_t/GFAL2Impl_t.py b/test/python/WMCore_t/Storage_t/Backends_t/GFAL2Impl_t.py
index 0ac59809e2..58a5e9fa84 100644
--- a/test/python/WMCore_t/Storage_t/Backends_t/GFAL2Impl_t.py
+++ b/test/python/WMCore_t/Storage_t/Backends_t/GFAL2Impl_t.py
@@ -14,11 +14,12 @@ def setUp(self):
 
     def testInit(self):
         testGFAL2Impl = GFAL2Impl()
+        # The default setup without a token
         removeCommand = "env -i JOBSTARTDIR=$JOBSTARTDIR bash -c " \
-                "'. $JOBSTARTDIR/startup_environment.sh; date; gfal-rm -t 600 {}'"
+                        "'. $JOBSTARTDIR/startup_environment.sh; date; gfal-rm -t 600 {}'"
         copyCommand = "env -i JOBSTARTDIR=$JOBSTARTDIR bash -c '" \
-              ". $JOBSTARTDIR/startup_environment.sh; date; gfal-copy -t 2400 -T 2400 -p " \
-              "-v --abort-on-failure {checksum} {options} {source} {destination}'"
+                    ". $JOBSTARTDIR/startup_environment.sh; date; gfal-copy -t 2400 -T 2400 -p " \
+                    "-v --abort-on-failure {checksum} {options} {source} {destination}'"
         self.assertEqual(removeCommand, testGFAL2Impl.removeCommand)
         self.assertEqual(copyCommand, testGFAL2Impl.copyCommand)
 
@@ -79,10 +80,23 @@ def testCreateRemoveFileCommand_removeCommand(self, mock_path):
     def testCreateStageOutCommand_stageIn(self, mock_createRemoveFileCommand):
         self.GFAL2Impl.stageIn = True
         mock_createRemoveFileCommand.return_value = "targetPFN2"
-        result = self.GFAL2Impl.createStageOutCommand("sourcePFN", "targetPFN")
+
+        # Call createStageOutCommand with auth_method='TOKEN'
+        result = self.GFAL2Impl.createStageOutCommand(
+            "sourcePFN", "targetPFN", auth_method='TOKEN'
+        )
+
+        # Generate the expected result with auth_method='TOKEN'
         expectedResult = self.getStageOutCommandResult(
-            self.getCopyCommandDict("-K adler32", "", "sourcePFN", "targetPFN"), "targetPFN2")
+            self.getCopyCommandDict("-K adler32", "", "sourcePFN", "targetPFN"),
+            "targetPFN2",
+            auth_method="TOKEN"
+        )
+
+        # Assert that the removeFileCommand was called correctly
         mock_createRemoveFileCommand.assert_called_with("targetPFN")
+
+        # Compare the expected and actual result
         self.assertEqual(expectedResult, result)
 
     @mock.patch('WMCore.Storage.Backends.GFAL2Impl.GFAL2Impl.createRemoveFileCommand')
@@ -94,20 +108,37 @@ def testCreateStageOutCommand_options(self, mock_createRemoveFileCommand):
         mock_createRemoveFileCommand.assert_called_with("file:targetPFN")
         self.assertEqual(expectedResult, result)
 
-    def getCopyCommandDict(self, checksum, options, source, destination):
-        copyCommandDict = {'checksum': '', 'options': '', 'source': '', 'destination': ''}
-        copyCommandDict['checksum'] = checksum
-        copyCommandDict['options'] = options
-        copyCommandDict['source'] = source
-        copyCommandDict['destination'] = destination
+    def getCopyCommandDict(self, checksum, options, source, destination, auth_method=None):
+        """
+        Generate a dictionary for the gfal-copy command, dynamically adjusting for auth_method.
+        """
+        copyCommandDict = {
+            'checksum': checksum,
+            'options': options,
+            'source': source,
+            'destination': destination
+        }
         return copyCommandDict
 
-    def getStageOutCommandResult(self, copyCommandDict, createRemoveFileCommandResult):
+    def getStageOutCommandResult(self, copyCommandDict, createRemoveFileCommandResult, auth_method=None):
+        """
+        Generate the expected result for the gfal-copy command, including dynamic adjustments for auth_method.
+        """
+        # Adjust the setup based on auth_method
+        if auth_method == "X509":
+            setups = "env -i X509_USER_PROXY=$X509_USER_PROXY JOBSTARTDIR=$JOBSTARTDIR bash -c '{}'"
+        elif auth_method == "TOKEN":
+            setups = "env -i BEARER_TOKEN=$(cat $BEARER_TOKEN_FILE) JOBSTARTDIR=$JOBSTARTDIR bash -c '{}'"
+        else:
+            setups = "env -i JOBSTARTDIR=$JOBSTARTDIR bash -c '{}'"
+
+        # Build the copy command dynamically
+        copyOpts = '-t 2400 -T 2400 -p -v --abort-on-failure {checksum} {options} {source} {destination}'
+        copyCommand = setups.format('. $JOBSTARTDIR/startup_environment.sh; date; gfal-copy ' + copyOpts)
+
+        # Construct the full result
         result = "#!/bin/bash\n"
-
-        copyCommand = self.copyCommand.format_map(copyCommandDict)
-        result += copyCommand
-
+        result += copyCommand.format_map(copyCommandDict)
         result += """
             EXIT_STATUS=$?
             echo "gfal-copy exit status: $EXIT_STATUS"
@@ -118,7 +149,7 @@ def getStageOutCommandResult(self, copyCommandDict, createRemoveFileCommandResul
             fi
             exit $EXIT_STATUS
             """.format(remove_command=createRemoveFileCommandResult)
-        
+
         return result
     
     @mock.patch('WMCore.Storage.Backends.GFAL2Impl.os.path')