Stabilize the DPU kernel panic and memory exhaustion tests

We should check the DPUs are offline before checking they are online
after the kernel panic or memory exhaustion.
Otherwise the check for DPU online could pass even before the DPUs are
rebooted and the later crirical services check will fail.

Signed-off-by: Cong Hou <congh@nvidia.com>

Author: congh-nvidia

tests/smartswitch/common/device_utils_dpu.py

@@ -225,6 +225,25 @@ def check_dpu_module_status(duthost, power_status, dpu_name):
             return False
 
 
+def check_dpus_module_status(duthost, dpu_list, power_status, wait_timeout=30):
+    """
+    Check module status of given DPU list
+    Args:
+        duthost : Host handle
+        dpu_list: List of DPUs to be checked
+        num_dpu_modules: number of dpu modules
+        power_status: status to be checked (on/off)
+        wait_timeout: timeout for the check
+    """
+    logging.info("Check module status of DPUs")
+    wait_interval = 10 if wait_timeout > 20 else wait_timeout // 3
+    for dpu_name in dpu_list:
+        pytest_assert(wait_until(wait_timeout, wait_interval, 0,
+                                 check_dpu_module_status, duthost,
+                                 power_status, dpu_name),
+                      f"DPU {dpu_name} is not {power_status}")
+
+
 def check_dpu_reboot_cause(duthost, dpu_name, reason):
     """
     Check reboot cause of all DPU modules

tests/smartswitch/platform_tests/test_reload_dpu.py

@@ -11,7 +11,7 @@
 from tests.common.reboot import reboot, REBOOT_TYPE_COLD, SONIC_SSH_PORT, SONIC_SSH_REGEX
 from tests.smartswitch.common.device_utils_dpu import check_dpu_link_and_status,\
     pre_test_check, post_test_switch_check, post_test_dpus_check,\
-    dpus_shutdown_and_check, dpus_startup_and_check,\
+    dpus_shutdown_and_check, dpus_startup_and_check, check_dpus_module_status,\
     num_dpu_modules, check_dpus_are_not_pingable, check_dpus_reboot_cause  # noqa: F401
 from tests.common.platform.device_utils import platform_api_conn, start_platform_api_service  # noqa: F401,F403
 from tests.smartswitch.common.reboot import perform_reboot
@@ -216,6 +216,9 @@ def test_dpu_status_post_dpu_kernel_panic(duthosts, dpuhosts,
 
         logging.info("Starting UP the DPUs")
         dpus_startup_and_check(duthost, dpu_on_list, num_dpu_modules)
+    else:
+        logging.info("Check DPUs are offline")
+        check_dpus_module_status(duthost, dpu_on_list, "off")
 
     logging.info("Executing post test dpu check")
     post_test_dpus_check(duthost, dpuhosts,
@@ -267,6 +270,9 @@ def test_dpu_check_post_dpu_mem_exhaustion(duthosts, dpuhosts,
 
         logging.info("Starting UP the DPUs")
         dpus_startup_and_check(duthost, dpu_on_list, num_dpu_modules)
+    else:
+        logging.info("Check DPUs are offline")
+        check_dpus_module_status(duthost, dpu_on_list, "off")
 
     logging.info("Executing post test dpu check")
     post_test_dpus_check(duthost, dpuhosts, dpu_on_list, ip_address_list,