From f31a6df098f091f52c325e702661cf7008483fb5 Mon Sep 17 00:00:00 2001
From: Gregory Boudreau <45526465+gregoryboudreau@users.noreply.github.com>
Date: Mon, 11 Mar 2024 10:43:21 -0700
Subject: [PATCH 1/2] psu key to use get_name api if available

---
 sonic-psud/scripts/psud | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/sonic-psud/scripts/psud b/sonic-psud/scripts/psud
index 3bcae4c76..8d8a25451 100644
--- a/sonic-psud/scripts/psud
+++ b/sonic-psud/scripts/psud
@@ -122,6 +122,11 @@ def _wrapper_get_psu_status(psu_index):
 #
 
 def get_psu_key(psu_index):
+    if platform_chassis is not None:
+        try:
+            return platform_chassis.get_psu(psu_index - 1).get_name()
+        except NotImplementedError:
+            pass
     return PSU_INFO_KEY_TEMPLATE.format(psu_index)
 
 

From 677642afa4443b98f53917b50dcbd6de65d0fe8b Mon Sep 17 00:00:00 2001
From: Gregory Boudreau <45526465+gregoryboudreau@users.noreply.github.com>
Date: Mon, 8 Jul 2024 04:17:39 -0700
Subject: [PATCH 2/2] modifications for tests

---
 sonic-psud/scripts/psud                 |  3 +++
 sonic-psud/tests/test_DaemonPsud.py     | 11 +++++------
 sonic-psud/tests/test_PsuChassisInfo.py | 11 +++++++++++
 sonic-psud/tests/test_psud.py           |  1 -
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/sonic-psud/scripts/psud b/sonic-psud/scripts/psud
index 8d8a25451..c27000d7c 100644
--- a/sonic-psud/scripts/psud
+++ b/sonic-psud/scripts/psud
@@ -127,6 +127,9 @@ def get_psu_key(psu_index):
             return platform_chassis.get_psu(psu_index - 1).get_name()
         except NotImplementedError:
             pass
+        except IndexError:
+            #some functionality is expectent on returning an expected key even if the psu object itself does not exist
+            pass
     return PSU_INFO_KEY_TEMPLATE.format(psu_index)
 
 
diff --git a/sonic-psud/tests/test_DaemonPsud.py b/sonic-psud/tests/test_DaemonPsud.py
index 05d7d5dda..ca866eca0 100644
--- a/sonic-psud/tests/test_DaemonPsud.py
+++ b/sonic-psud/tests/test_DaemonPsud.py
@@ -517,20 +517,19 @@ def test_update_psu_entity_info(self):
         assert daemon_psud.log_warning.call_count == 1
         daemon_psud.log_warning.assert_called_with("Failed to update PSU data - Test message")
 
-    @mock.patch('psud.try_get', mock.MagicMock(return_value=0))
     def test_update_single_psu_entity_info(self):
-        mock_psu1 = MockPsu("PSU 1", 0, True, True)
-
+        #creating psu object in slot not used to allow for name specific check
+        mock_psu1 = MockPsu("PSU 3", 2, True, True)
         expected_fvp = psud.swsscommon.FieldValuePairs(
-            [('position_in_parent', '0'),
+            [('position_in_parent', '2'),
              ('parent_name', psud.CHASSIS_INFO_KEY),
              ])
 
         daemon_psud = psud.DaemonPsud(SYSLOG_IDENTIFIER)
         daemon_psud.phy_entity_tbl = mock.MagicMock()
 
-        daemon_psud._update_single_psu_entity_info(0, mock_psu1)
-        daemon_psud.phy_entity_tbl.set.assert_called_with('PSU 0', expected_fvp)
+        daemon_psud._update_single_psu_entity_info(3, mock_psu1)
+        daemon_psud.phy_entity_tbl.set.assert_called_with('PSU 3', expected_fvp)
 
     @mock.patch('psud.datetime')
     def test_update_psu_fan_data(self, mock_datetime):
diff --git a/sonic-psud/tests/test_PsuChassisInfo.py b/sonic-psud/tests/test_PsuChassisInfo.py
index f40836268..196ef2d48 100644
--- a/sonic-psud/tests/test_PsuChassisInfo.py
+++ b/sonic-psud/tests/test_PsuChassisInfo.py
@@ -290,9 +290,20 @@ def test_first_run(self):
         assert psud.Psu.get_status_master_led() == MockPsu.STATUS_LED_COLOR_GREEN
 
     def test_get_psu_key(self):
+        
+        #mock test for psu get_name() nonimplementation
+        psud.platform_chassis = None
         assert psud.get_psu_key(0) == psud.PSU_INFO_KEY_TEMPLATE.format(0)
         assert psud.get_psu_key(1) == psud.PSU_INFO_KEY_TEMPLATE.format(1)
 
+        #create psu objects to test getting name for
+        mock_psu1 = MockPsu("PSU 1", 0, True, True)
+        mock_psu2 = MockPsu("PSU 2", 1, True, True)
+        psud.platform_chassis = MockChassis()
+        psud.platform_chassis._psu_list = [mock_psu1, mock_psu2]
+        assert psud.get_psu_key(1) == psud.PSU_INFO_KEY_TEMPLATE.format(1)
+        assert psud.get_psu_key(2) == psud.PSU_INFO_KEY_TEMPLATE.format(2)
+
     def test_try_get(self):
         # Test a proper, working callback
         GOOD_CALLBACK_RETURN_VALUE = "This is a test"
diff --git a/sonic-psud/tests/test_psud.py b/sonic-psud/tests/test_psud.py
index b95a4986a..463ee9edd 100644
--- a/sonic-psud/tests/test_psud.py
+++ b/sonic-psud/tests/test_psud.py
@@ -95,7 +95,6 @@ def test_wrapper_get_psu_status():
     assert psud.platform_psuutil.get_psu_status.call_count == 2
     psud.platform_psuutil.get_psu_status.assert_called_with(1)
 
-
 def test_log_on_status_changed():
     normal_log = "Normal log message"
     abnormal_log = "Abnormal log message"