From 10a2cc41e6ad25057e071221d5f9d75329f4e4ee Mon Sep 17 00:00:00 2001
From: Sonic Build Admin <sonicbld@microsoft.com>
Date: Wed, 29 Jan 2025 23:07:57 +0000
Subject: [PATCH] Add T2 support to test_power_off_reboot.py

<!--
Please make sure you've read and understood our contributing guidelines;
https://github.com/sonic-net/SONiC/blob/gh-pages/CONTRIBUTING.md

Please provide following information to help code review process a bit easier:
-->
### Description of PR
<!--
- Please include a summary of the change and which issue is fixed.
- Please also include relevant motivation and context. Where should reviewer start? background context?
- List any dependencies that are required for this change.
-->

Summary:
Fixes #16279, update platform_tests/power_off_reboot.py to support T2 devices

### Type of change

<!--
- Fill x for your type of change.
- e.g.
- [x] Bug fix
-->

- [x] Bug fix
- [ ] Testbed and Framework(new/improvement)
- [ ] Test case(new/improvement)

### Back port request
- [ ] 202012
- [ ] 202205
- [ ] 202305
- [ ] 202311
- [x] 202405

### Approach
#### What is the motivation for this PR?
Test failures on T2 devices
#### How did you do it?
Change check for power on each PDU to check for power on sets of PDUs (by PSU), add duthosts argument to `reboot_and_check` function to support sup cards
#### How did you verify/test it?
Perform regression test on T1 https://elastictest.org/scheduler/testplan/677b439de1be6860ea024869?testcase=platform_tests%2Ftest_power_off_reboot.py&type=console

Perform test on T2 to ensure these failures are gone (still failure in #16289)
#### Any platform specific information?
N/A
#### Supported testbed topology if it's a new test case?
N/A
### Documentation
<!--
(If it's a new feature, new test case)
Did you update documentation/Wiki relevant to your implementation?
Link to the wiki page?
-->
---
 tests/platform_tests/test_power_off_reboot.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tests/platform_tests/test_power_off_reboot.py b/tests/platform_tests/test_power_off_reboot.py
index 0f4fd31f79..d67c423d05 100644
--- a/tests/platform_tests/test_power_off_reboot.py
+++ b/tests/platform_tests/test_power_off_reboot.py
@@ -4,7 +4,8 @@
 from tests.common.reboot import wait_for_startup, REBOOT_TYPE_POWEROFF
 from tests.common.platform.processes_utils import wait_critical_processes, check_critical_processes
 from tests.common.helpers.assertions import pytest_assert
-from tests.platform_tests.test_reboot import check_interfaces_and_services,\
+from tests.common.helpers.psu_helpers import get_grouped_pdus_by_psu
+from tests.platform_tests.test_reboot import check_interfaces_and_services, \
     reboot_and_check
 from tests.common.utilities import get_plt_reboot_ctrl
 
@@ -93,9 +94,10 @@ def test_power_off_reboot(duthosts, localhost, enum_supervisor_dut_hostname, con
         # Following is to accomodate for chassis, when no '--power_off_delay' option is given on pipeline run
         power_off_delay = 60
     all_outlets = pdu_ctrl.get_outlet_status()
-    # If PDU supports returning output_watts, making sure that all outlets has power.
-    no_power = [item for item in all_outlets if int(item.get('output_watts', '1')) == 0]
-    pytest_assert(not no_power, "Not all outlets have power output: {}".format(no_power))
+    # If PDU supports returning output_watts, making sure that all PSUs has power.
+    psu_to_pdus = get_grouped_pdus_by_psu(pdu_ctrl)
+    for psu, pdus in psu_to_pdus.items():
+        pytest_assert(any(pdu["output_watts"] != 0 for pdu in pdus), "Not all PSUs are getting power")
 
     # Purpose of this list is to control sequence of turning on PSUs in power off testing.
     # If there are 2 PSUs, then 3 scenarios would be covered:
@@ -104,7 +106,7 @@ def test_power_off_reboot(duthosts, localhost, enum_supervisor_dut_hostname, con
     # 3. Turn off all PSUs, turn on one of the PSU, then turn on the other PSU, then check.
     power_on_seq_list = []
     if all_outlets:
-        power_on_seq_list = [[item] for item in all_outlets]
+        power_on_seq_list = [pdus for pdus in psu_to_pdus.values()]
         power_on_seq_list.append(all_outlets)
 
     logging.info("Got all power on sequences {}".format(power_on_seq_list))
@@ -118,8 +120,9 @@ def test_power_off_reboot(duthosts, localhost, enum_supervisor_dut_hostname, con
             poweroff_reboot_kwargs["power_on_seq"] = power_on_seq
             poweroff_reboot_kwargs["delay_time"] = power_off_delay
             reboot_and_check(localhost, duthost, conn_graph_facts.get("device_conn", {}).get(duthost.hostname, {}),
-                             xcvr_skip_list, REBOOT_TYPE_POWEROFF,
-                             _power_off_reboot_helper, poweroff_reboot_kwargs)
+                             xcvr_skip_list, REBOOT_TYPE_POWEROFF, _power_off_reboot_helper,
+                             poweroff_reboot_kwargs, duthosts)
+
     except Exception as e:
         logging.debug("Restore power after test failure")
         for outlet in all_outlets: