Add unit tests for reboot_helper.py

sonic-net · Oct 25, 2024 · b2dcd63 · b2dcd63
1 parent 01cf117
commit b2dcd63
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 5 deletions.
diff --git a/scripts/reboot b/scripts/reboot
@@ -11,8 +11,6 @@ PLATFORM_UPDATE_REBOOT_CAUSE="platform_update_reboot_cause"
 REBOOT_CAUSE_FILE="/host/reboot-cause/reboot-cause.txt"
 PLATFORM_REBOOT_PRE_CHECK="platform_reboot_pre_check"
 REBOOT_TIME=$(date)
-PLATFORM_JSON_FILE="platform.json"
-PLATFORM_JSON_PATH="${DEVPATH}/${PLATFORM}/${PLATFORM_JSON_FILE}"
 
 # Reboot immediately if we run the kdump capture kernel
 VMCORE_FILE=/proc/vmcore
@@ -39,6 +37,8 @@ EXIT_NEXT_IMAGE_NOT_EXISTS=4
 EXIT_SONIC_INSTALLER_VERIFY_REBOOT=21
 EXIT_PLATFORM_FW_AU_FAILURE=22
 PLATFORM_FWUTIL_AU_REBOOT_HANDLE="platform_fw_au_reboot_handle"
+PLATFORM_JSON_FILE="platform.json"
+PLATFORM_JSON_PATH="${DEVPATH}/${PLATFORM}/${PLATFORM_JSON_FILE}"
 REBOOT_SCRIPT_NAME=$(basename $0)
 REBOOT_TYPE="${REBOOT_SCRIPT_NAME}"
 TAG_LATEST=no
@@ -213,7 +213,7 @@ function get_reboot_status()
 {
     local dpu_ip=$1
     local port=$2
-    reboot_status=$(gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -rpc RebootStatus)
+    reboot_status=$(docker exec -i gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -rpc RebootStatus)
     if [ $? -ne 0 ] || [ -z "$reboot_status" ]; then
         echo "Error: Failed to send reboot status command to DPU ${DPU_NAME}"
         exit ${EXIT_ERROR}
@@ -259,7 +259,7 @@ function reboot_dpu_module()
     fi
 
     # Issue GNOI client command to reboot the DPU
-    gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -rpc Reboot -jsonin '{"method":3}'
+    docker exec -i gnmi gnoi_client -target ${dpu_ip}:${port} -logtostderr -insecure -rpc Reboot -jsonin '{"method":3}'
     if [ $? -ne 0 ]; then
         echo "Error: Failed to send reboot command to DPU ${DPU_NAME}"
         exit ${EXIT_ERROR}

diff --git a/scripts/reboot_helper.py b/scripts/reboot_helper.py
@@ -4,9 +4,9 @@
 #
 # Utility helper for reboot within SONiC
 
-import sonic_platform
 import sys
 import syslog
+import sonic_platform
 
 chk_log_level = syslog.LOG_ERR
 
@@ -58,6 +58,9 @@ def reboot_module(module_name):
         log_err("Failed to load platform chassis")
         return False
 
+    if not platform_chassis.is_smartswitch():
+        return False
+
     # Iterate over the modules to find the one with the specified name
     try:
         # Use get_all_modules to retrieve all modules on the chassis

diff --git a/tests/test_reboot_helper.py b/tests/test_reboot_helper.py
@@ -0,0 +1,111 @@
+import os
+import sys
+import unittest
+from unittest.mock import patch, MagicMock
+
+import pytest
+from utilities_common.general import load_module_from_source
+
+test_path = os.path.dirname(os.path.abspath(__file__))
+modules_path = os.path.dirname(test_path)
+scripts_path = os.path.join(modules_path, "scripts")
+sys.path.insert(0, modules_path)
+
+reboot_helper_path = os.path.join(scripts_path, 'reboot_helper.py')
+reboot_helper = load_module_from_source('reboot_helper', reboot_helper_path)
+
+class TestRebootHelper(unittest.TestCase):
+
+    @patch('reboot_helper.syslog.syslog')
+    def test_log_err(self, mock_syslog):
+        reboot_helper.log_err('Test error message')
+        mock_syslog.assert_called_with(reboot_helper.syslog.LOG_ERR, 'Test error message')
+
+    @patch('reboot_helper.syslog.syslog')
+    def test_log_info(self, mock_syslog):
+        reboot_helper.log_info('Test info message')
+        mock_syslog.assert_called_with(reboot_helper.syslog.LOG_INFO, 'Test info message')
+
+    @patch('reboot_helper.syslog.syslog')
+    def test_log_debug(self, mock_syslog):
+        reboot_helper.log_debug('Test debug message')
+        mock_syslog.assert_called_with(reboot_helper.syslog.LOG_DEBUG, 'Test debug message')
+
+    @patch('reboot_helper.sonic_platform.platform.Platform')
+    def test_load_platform_chassis_success(self, mock_platform):
+        mock_chassis = MagicMock()
+        mock_platform.return_value.get_chassis.return_value = mock_chassis
+        self.assertTrue(reboot_helper.load_platform_chassis())
+        self.assertEqual(reboot_helper.platform_chassis, mock_chassis)
+
+    @patch('reboot_helper.sonic_platform.platform.Platform')
+    def test_load_platform_chassis_failure(self, mock_platform):
+        mock_platform.return_value.get_chassis.side_effect = Exception('Load failed')
+        self.assertFalse(reboot_helper.load_platform_chassis())
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_reboot_module_success(self, mock_get_all_modules, mock_load_platform_chassis):
+        mock_module = MagicMock()
+        mock_module.get_name.return_value = 'test_module'
+        mock_get_all_modules.return_value = [mock_module]
+        self.assertTrue(reboot_helper.reboot_module('test_module'))
+        mock_module.reboot.assert_called_once()
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_reboot_module_not_implemented_failure(self, mock_get_all_modules, mock_load_platform_chassis):
+        mock_module = MagicMock()
+        mock_module.get_name.return_value = 'test_module'
+        mock_module.reboot.side_effect = NotImplementedError
+        mock_get_all_modules.return_value = [mock_module]
+        self.assertFalse(reboot_helper.reboot_module('test_module'))
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_reboot_module_general_exception(self, mock_get_all_modules, mock_load_platform_chassis):
+        mock_module = MagicMock()
+        mock_module.get_name.return_value = 'test_module'
+        mock_module.reboot.side_effect = Exception('Reboot failed')
+        mock_get_all_modules.return_value = [mock_module]
+        self.assertFalse(reboot_helper.reboot_module('test_module'))
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules', side_effect=Exception('Failed to get modules'))
+    def test_reboot_module_exception(self, mock_get_all_modules, mock_load_platform_chassis):
+        self.assertFalse(reboot_helper.reboot_module('test_module'))
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.is_smartswitch', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_is_dpu_success(self, mock_get_all_modules, mock_is_smartswitch, mock_load_platform_chassis):
+        mock_module = MagicMock()
+        mock_module.get_name.return_value = 'DPU test_module'
+        mock_get_all_modules.return_value = [mock_module]
+        self.assertTrue(reboot_helper.is_dpu())
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.is_smartswitch', return_value=False)
+    def test_is_dpu_failure(self, mock_is_smartswitch, mock_load_platform_chassis):
+        self.assertFalse(reboot_helper.is_dpu())
+
+    @patch('reboot_helper.load_platform_chassis', return_value=False)
+    def test_is_dpu_load_platform_chassis_failure(self, mock_load_platform_chassis):
+        self.assertFalse(reboot_helper.is_dpu())
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_reboot_module_not_found(self, mock_get_all_modules, mock_load_platform_chassis):
+        mock_get_all_modules.return_value = []
+        self.assertFalse(reboot_helper.reboot_module('non_existent_module'))
+
+    @patch('reboot_helper.load_platform_chassis', return_value=True)
+    @patch('reboot_helper.platform_chassis.get_all_modules')
+    def test_reboot_module_name_mismatch(self, mock_get_all_modules, mock_load_platform_chassis):
+        mock_module = MagicMock()
+        mock_module.get_name.return_value = 'another_module'
+        mock_get_all_modules.return_value = [mock_module]
+        self.assertFalse(reboot_helper.reboot_module('test_module'))
+
+if __name__ == '__main__':
+    unittest.main()