patch.patch

From ee0990275bc1939f6017c714fcb0b4fb85477a2f Mon Sep 17 00:00:00 2001
From: rdimaio <riccardodimaio11@gmail.com>
Date: Fri, 25 Oct 2024 18:52:39 +0200
Subject: [PATCH] Transfers: Add TestCollocationPlugin to FTS3 plugin tests

---
 tests/test_transfer_plugins.py | 159 ++++++++++++++++++---------------
 1 file changed, 86 insertions(+), 73 deletions(-)

diff --git a/tests/test_transfer_plugins.py b/tests/test_transfer_plugins.py
index 1de5be0257..0c40d45b04 100644
--- a/tests/test_transfer_plugins.py
+++ b/tests/test_transfer_plugins.py
@@ -25,6 +25,7 @@
 from rucio.core.transfer import ProtocolFactory, build_transfer_paths
 from rucio.db.sqla.session import get_session
 from rucio.transfertool.fts3 import FTS3Transfertool, build_job_params
+from rucio.transfertool.fts3_plugins import FTS3TapeMetadataPlugin
 
 mock_session = get_session()
 
@@ -154,89 +155,101 @@ def test_activity_missing(file_config_mock, did_factory, rse_factory, root_accou
     assert expected_scheduling_hints == generated_scheduling_hints
 
 
-@pytest.mark.parametrize("file_config_mock", [
-    {
-        "overrides": [
-            ("transfers", "fts3tape_metadata_plugins", "test")
-        ]
-    }
-], indirect=True)
-def test_collocation_hints(file_config_mock, did_factory, rse_factory, root_account):
-    """For a mock collocation algorithm, it can produce the 4 levels of hints required for each did"""
+class TestCollocationHints:
+    class TestCollocationPlugin(FTS3TapeMetadataPlugin):
+        def __init__(self) -> None:
+            self.register(
+                'test',
+                func=lambda x: self._test_collocation(x))
+            super().__init__('test')
 
-    mock_did = did_factory.random_file_did()
-    transfer_path = _make_transfer_path(mock_did, rse_factory, root_account)
+        def _test_collocation(self, hints: dict[str, str]) -> dict[str, dict]:
+            return {"collocation_hints": {"0": "", "1": "", "2": "", "3": ""}}
 
-    # Mock Transfer Tool
-    fts3_tool = FTS3Transfertool(TEST_FTS_HOST)
+    TestCollocationPlugin()
 
-    job_params = build_job_params(
-        transfer_path=transfer_path,
-        bring_online=None,
-        default_lifetime=None,
-        archive_timeout_override=None,
-        max_time_in_queue=None,
-        logger=logging.log,
-    )
-
-    # Get the job params used for each transfer
-    job_params = fts3_tool._file_from_transfer(transfer_path[0], job_params)
-
-    expected_collocation_hints = {
-        "collocation_hints": {
-            "0": "",
-            "1": "",
-            "2": "",
-            "3": "",
+    @pytest.mark.parametrize("file_config_mock", [
+        {
+            "overrides": [
+                ("transfers", "fts3tape_metadata_plugins", "test")
+            ]
         }
-    }
-
-    assert "archive_metadata" in job_params
-    generated_collocation_hints = job_params["archive_metadata"]["collocation_hints"]
-
-    assert (
-        expected_collocation_hints["collocation_hints"] == generated_collocation_hints
-    )
-
-
-@pytest.mark.parametrize("file_config_mock", [
-    {
-        "overrides": [
-            ("transfers", "fts3tape_metadata_plugins", "activity, test")
-        ]
-    }
-], indirect=True)
-def test_multiple_plugin_concat(file_config_mock, did_factory, rse_factory, root_account):
-    """When multiple plugins are used (like priority and collocation), both logics are applied"""
+    ], indirect=True)
+    def test_collocation_hints(self, file_config_mock, did_factory, rse_factory, root_account):
+        """For a mock collocation algorithm, it can produce the 4 levels of hints required for each did"""
+
+        mock_did = did_factory.random_file_did()
+        transfer_path = _make_transfer_path(mock_did, rse_factory, root_account)
+
+        # Mock Transfer Tool
+        fts3_tool = FTS3Transfertool(TEST_FTS_HOST)
+
+        job_params = build_job_params(
+            transfer_path=transfer_path,
+            bring_online=None,
+            default_lifetime=None,
+            archive_timeout_override=None,
+            max_time_in_queue=None,
+            logger=logging.log,
+        )
+
+        # Get the job params used for each transfer
+        job_params = fts3_tool._file_from_transfer(transfer_path[0], job_params)
 
-    mock_did = did_factory.random_file_did()
-    transfer_path = _make_transfer_path(mock_did, rse_factory, root_account)
+        expected_collocation_hints = {
+            "collocation_hints": {
+                "0": "",
+                "1": "",
+                "2": "",
+                "3": "",
+            }
+        }
 
-    # Mock Transfer Tool
-    fts3_tool = FTS3Transfertool(TEST_FTS_HOST)
+        assert "archive_metadata" in job_params
+        generated_collocation_hints = job_params["archive_metadata"]["collocation_hints"]
 
-    job_params = build_job_params(
-        transfer_path=transfer_path,
-        bring_online=None,
-        default_lifetime=None,
-        archive_timeout_override=None,
-        max_time_in_queue=None,
-        logger=logging.log,
-    )
+        assert (
+            expected_collocation_hints["collocation_hints"] == generated_collocation_hints
+        )
 
-    # Get the job params used for each transfer
-    job_params = fts3_tool._file_from_transfer(transfer_path[0], job_params)
-    expected_hints = {
-        "scheduling_hints": {"priority": "20"},
-        "collocation_hints": {"0": "", "1": "", "2": "", "3": ""},
-    }
-    assert "archive_metadata" in job_params
+    @pytest.mark.parametrize("file_config_mock", [
+        {
+            "overrides": [
+                ("transfers", "fts3tape_metadata_plugins", "activity, test")
+            ]
+        }
+    ], indirect=True)
+    def test_multiple_plugin_concat(self, file_config_mock, did_factory, rse_factory, root_account):
+        """When multiple plugins are used (like priority and collocation), both logics are applied"""
+
+        mock_did = did_factory.random_file_did()
+        transfer_path = _make_transfer_path(mock_did, rse_factory, root_account)
+
+        # Mock Transfer Tool
+        fts3_tool = FTS3Transfertool(TEST_FTS_HOST)
+
+        job_params = build_job_params(
+            transfer_path=transfer_path,
+            bring_online=None,
+            default_lifetime=None,
+            archive_timeout_override=None,
+            max_time_in_queue=None,
+            logger=logging.log,
+        )
+
+        # Get the job params used for each transfer
+        job_params = fts3_tool._file_from_transfer(transfer_path[0], job_params)
+        expected_hints = {
+            "scheduling_hints": {"priority": "20"},
+            "collocation_hints": {"0": "", "1": "", "2": "", "3": ""},
+        }
+        assert "archive_metadata" in job_params
 
-    generated_collocation_hints = job_params["archive_metadata"]["collocation_hints"]
-    assert expected_hints["collocation_hints"] == generated_collocation_hints
+        generated_collocation_hints = job_params["archive_metadata"]["collocation_hints"]
+        assert expected_hints["collocation_hints"] == generated_collocation_hints
 
-    generated_scheduling_hints = job_params["archive_metadata"]["scheduling_hints"]
-    assert expected_hints["scheduling_hints"] == generated_scheduling_hints
+        generated_scheduling_hints = job_params["archive_metadata"]["scheduling_hints"]
+        assert expected_hints["scheduling_hints"] == generated_scheduling_hints
 
 
 @pytest.mark.parametrize("file_config_mock", [