diff --git a/host_modules/systemd_service.py b/host_modules/systemd_service.py index 19a6dd60..974af1e0 100644 --- a/host_modules/systemd_service.py +++ b/host_modules/systemd_service.py @@ -1,5 +1,6 @@ """Systemd service handler""" +from enum import Enum from host_modules import host_service import subprocess @@ -7,6 +8,11 @@ ALLOWED_SERVICES = ['snmp', 'swss', 'dhcp_relay', 'radv', 'restapi', 'lldp', 'sshd', 'pmon', 'rsyslog', 'telemetry'] EXIT_FAILURE = 1 +# Define an Enum for Reboot Methods which are defined as in +# https://github.com/openconfig/gnoi/blob/main/system/system.pb.go#L27 +class RebootMethod(Enum): + COLD = 1 + HALT = 3 class SystemdService(host_service.HostModule): """ @@ -48,3 +54,19 @@ def stop_service(self, service): if result.returncode: msg = result.stderr.decode() return result.returncode, msg + + @host_service.method(host_service.bus_name(MOD_NAME), in_signature='i', out_signature='is') + def execute_reboot(self, rebootmethod): + if rebootmethod == RebootMethod.COLD: + cmd = ['/usr/local/bin/reboot'] + elif rebootmethod == RebootMethod.HALT: + cmd = ['/usr/local/bin/reboot','-p'] + else: + return EXIT_FAILURE, "{}: Invalid reboot method: {}".format(MOD_NAME, rebootmethod) + + result = subprocess.run(cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + msg = '' + if result.returncode: + msg = result.stderr.decode() + + return result.returncode, msg diff --git a/tests/host_modules/systemd_service_test.py b/tests/host_modules/systemd_service_test.py index d8d5083f..54961498 100644 --- a/tests/host_modules/systemd_service_test.py +++ b/tests/host_modules/systemd_service_test.py @@ -84,3 +84,57 @@ def test_service_stop_empty(self, MockInit, MockBusName, MockSystemBus): ret, msg = systemd_service_stub.stop_service(service) assert ret == 1 assert "stop_service called with no service specified" in msg + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + def test_execute_reboot_cold(self, MockInit, MockBusName, MockSystemBus): + # Mock subprocess.run + with mock.patch("subprocess.run") as mock_run: + # Mock the result of subprocess.run + res_mock = mock.Mock() + test_ret = 0 + test_msg = b"Succeeded" + res_mock.configure_mock(returncode=test_ret, stderr=test_msg) + mock_run.return_value = res_mock + + method = systemd_service.RebootMethod.COLD + systemd_service_stub = systemd_service.SystemdService(systemd_service.MOD_NAME) + + # Execute the reboot method + ret, msg = systemd_service_stub.execute_reboot(method) + + # Assert the correct command was called + call_args = mock_run.call_args[0][0] + assert "/usr/local/bin/reboot" in call_args, f"Expected 'reboot' command, but got: {call_args}" + + # Assert the return values are correct + assert ret == test_ret, f"Expected return code {test_ret}, got {ret}" + assert msg == "", f"Expected return message '', got {msg}" + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + def test_execute_reboot_halt(self, MockInit, MockBusName, MockSystemBus): + # Mock subprocess.run + with mock.patch("subprocess.run") as mock_run: + # Mock the result of subprocess.run + res_mock = mock.Mock() + test_ret = 0 + test_msg = b"Succeeded" + res_mock.configure_mock(returncode=test_ret, stderr=test_msg) + mock_run.return_value = res_mock + + method = systemd_service.RebootMethod.HALT + systemd_service_stub = systemd_service.SystemdService(systemd_service.MOD_NAME) + + # Execute the reboot method + ret, msg = systemd_service_stub.execute_reboot(method) + + # Assert the correct command was called + call_args = mock_run.call_args[0][0] + assert "/usr/local/bin/reboot" in call_args, f"Expected 'reboot' command, but got: {call_args}" + + # Assert the return values are correct + assert ret == test_ret, f"Expected return code {test_ret}, got {ret}" + assert msg == "", f"Expected return message '', got {msg}"