From 9bd61362c1d0d4e9fa34d0026c960cea6090bd72 Mon Sep 17 00:00:00 2001
From: aldbr <aldbr@outlook.com>
Date: Tue, 10 Mar 2026 11:09:03 +0100
Subject: [PATCH] fix(rss): restore CEs in getSiteElements

---
 .../Utilities/CSHelpers.py                    | 12 +++-
 .../Utilities/test/Test_CSHelpers.py          | 65 +++++++++++++++++++
 .../Utilities/test/__init__.py                |  0
 3 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 src/DIRAC/ResourceStatusSystem/Utilities/test/Test_CSHelpers.py
 create mode 100644 src/DIRAC/ResourceStatusSystem/Utilities/test/__init__.py

diff --git a/src/DIRAC/ResourceStatusSystem/Utilities/CSHelpers.py b/src/DIRAC/ResourceStatusSystem/Utilities/CSHelpers.py
index 985dc099269..de086219a0b 100644
--- a/src/DIRAC/ResourceStatusSystem/Utilities/CSHelpers.py
+++ b/src/DIRAC/ResourceStatusSystem/Utilities/CSHelpers.py
@@ -4,6 +4,7 @@
 """
 
 from DIRAC import S_OK, gConfig, gLogger
+from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getCESiteMapping
 from DIRAC.Core.Utilities.SiteSEMapping import getSEParameters
 from DIRAC.DataManagementSystem.Utilities.DMSHelpers import DMSHelpers
 
@@ -75,12 +76,17 @@ def getFileCatalogs():
 
 def getSiteElements(siteName):
     """
-    Gets all the storage elements for a given site
+    Gets all the storage and computing elements for a given site
     """
 
     res = DMSHelpers().getSiteSEMapping()
     if not res["OK"]:
         return res
-    resources = res["Value"][1].get(siteName, [])
+    resources = list(res["Value"][1].get(siteName, []))
 
-    return S_OK(list(resources))
+    res = getCESiteMapping()
+    if not res["OK"]:
+        return res
+    resources.extend(ce for ce, site in res["Value"].items() if site == siteName)
+
+    return S_OK(resources)
diff --git a/src/DIRAC/ResourceStatusSystem/Utilities/test/Test_CSHelpers.py b/src/DIRAC/ResourceStatusSystem/Utilities/test/Test_CSHelpers.py
new file mode 100644
index 00000000000..5911a24c33a
--- /dev/null
+++ b/src/DIRAC/ResourceStatusSystem/Utilities/test/Test_CSHelpers.py
@@ -0,0 +1,65 @@
+"""Tests for CSHelpers module"""
+
+from unittest.mock import patch, MagicMock
+
+from DIRAC import S_OK, S_ERROR
+from DIRAC.ResourceStatusSystem.Utilities.CSHelpers import getSiteElements
+
+
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.getCESiteMapping")
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.DMSHelpers")
+def test_getSiteElements_returns_ses_and_ces(mock_dms_helpers, mock_get_ce_site_mapping):
+    """Test that getSiteElements returns both StorageElements and ComputingElements"""
+    mock_dms_instance = MagicMock()
+    mock_dms_instance.getSiteSEMapping.return_value = S_OK(({}, {"LCG.CERN.cern": ["CERN-SE1", "CERN-SE2"]}, {}))
+    mock_dms_helpers.return_value = mock_dms_instance
+
+    mock_get_ce_site_mapping.return_value = S_OK(
+        {"ce1.cern.ch": "LCG.CERN.cern", "ce2.cern.ch": "LCG.CERN.cern", "ce3.other.ch": "LCG.Other.cern"}
+    )
+
+    result = getSiteElements("LCG.CERN.cern")
+    assert result["OK"]
+    assert sorted(result["Value"]) == sorted(["CERN-SE1", "CERN-SE2", "ce1.cern.ch", "ce2.cern.ch"])
+
+
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.getCESiteMapping")
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.DMSHelpers")
+def test_getSiteElements_site_with_only_ces(mock_dms_helpers, mock_get_ce_site_mapping):
+    """Test that a site with only CEs (no SEs) still returns elements"""
+    mock_dms_instance = MagicMock()
+    mock_dms_instance.getSiteSEMapping.return_value = S_OK(({}, {}, {}))
+    mock_dms_helpers.return_value = mock_dms_instance
+
+    mock_get_ce_site_mapping.return_value = S_OK({"ce1.krakow.pl": "LCG.Krakow.pl"})
+
+    result = getSiteElements("LCG.Krakow.pl")
+    assert result["OK"]
+    assert result["Value"] == ["ce1.krakow.pl"]
+
+
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.getCESiteMapping")
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.DMSHelpers")
+def test_getSiteElements_ce_mapping_fails(mock_dms_helpers, mock_get_ce_site_mapping):
+    """Test that if CE mapping fails, error is propagated"""
+    mock_dms_instance = MagicMock()
+    mock_dms_instance.getSiteSEMapping.return_value = S_OK(({}, {"LCG.CERN.cern": ["CERN-SE1"]}, {}))
+    mock_dms_helpers.return_value = mock_dms_instance
+
+    mock_get_ce_site_mapping.return_value = S_ERROR("Failed to get CE mapping")
+
+    result = getSiteElements("LCG.CERN.cern")
+    assert not result["OK"]
+
+
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.getCESiteMapping")
+@patch("DIRAC.ResourceStatusSystem.Utilities.CSHelpers.DMSHelpers")
+def test_getSiteElements_se_mapping_fails(mock_dms_helpers, mock_get_ce_site_mapping):
+    """Test that if SE mapping fails, error is propagated"""
+    mock_dms_instance = MagicMock()
+    mock_dms_instance.getSiteSEMapping.return_value = S_ERROR("Failed to get SE mapping")
+    mock_dms_helpers.return_value = mock_dms_instance
+
+    result = getSiteElements("LCG.CERN.cern")
+    assert not result["OK"]
+    mock_get_ce_site_mapping.assert_not_called()
diff --git a/src/DIRAC/ResourceStatusSystem/Utilities/test/__init__.py b/src/DIRAC/ResourceStatusSystem/Utilities/test/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d