Skip to content

Commit

Permalink
Add unit tests for reboot_helper.py
Browse files Browse the repository at this point in the history
  • Loading branch information
vvolam committed Oct 25, 2024
1 parent 01cf117 commit b2dcd63
Show file tree
Hide file tree
Showing 3 changed files with 119 additions and 5 deletions.
8 changes: 4 additions & 4 deletions scripts/reboot
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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}
Expand Down
5 changes: 4 additions & 1 deletion scripts/reboot_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
111 changes: 111 additions & 0 deletions tests/test_reboot_helper.py
Original file line number Diff line number Diff line change
@@ -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()

0 comments on commit b2dcd63

Please sign in to comment.