diff --git a/host_modules/image_service.py b/host_modules/image_service.py index c52c6de9..cfe2c786 100644 --- a/host_modules/image_service.py +++ b/host_modules/image_service.py @@ -160,6 +160,27 @@ def list_images(self): logger.error(msg) return e.returncode, msg + @host_service.method( + host_service.bus_name(MOD_NAME), in_signature="s", out_signature="is" + ) + def set_next_boot(self, image): + """ + Set the image to be used for the next boot. + + Args: + image: The name of the image to set for the next boot. + """ + logger.info("Setting the next boot image to {}".format(image)) + cmd = ["/usr/local/bin/sonic-installer", "set-next-boot", image] + result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + msg = "Boot image set to {}".format(image) + logger.info(msg) + if result.returncode: + logger.error("Failed to set next boot image: {}".format(result.stderr.decode())) + msg = result.stderr.decode() + return result.returncode, msg + + def _parse_sonic_installer_list(self, output): """ Parse the output of the sonic-installer list command. diff --git a/tests/host_modules/image_service_test.py b/tests/host_modules/image_service_test.py index bc941778..686971f2 100644 --- a/tests/host_modules/image_service_test.py +++ b/tests/host_modules/image_service_test.py @@ -568,3 +568,59 @@ def test_list_images_failed(self, mock_check_output, MockInit, MockBusName, Mock stderr=subprocess.STDOUT, ) + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.run") + def test_image_set_next_boot_success(self, mock_run, MockInit, MockBusName, MockSystemBus): + """ + Test that the `set_next_boot` method successfully sets the next boot image. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image = "sonic_image" + mock_result = mock.Mock() + mock_result.returncode = 0 + mock_result.stderr = b"" + mock_run.return_value = mock_result + + # Act + rc, msg = image_service.set_next_boot(image) + + # Assert + assert rc == 0, "wrong return value" + assert image in msg, "message should contain the name of the new image" + mock_run.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "set-next-boot", image], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.run") + def test_image_set_next_boot_fail_not_exists(self, mock_run, MockInit, MockBusName, MockSystemBus): + """ + Test that the `set_next_boot` method fails when the image does not exist. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + image = "nonexistent_image" + mock_result = mock.Mock() + mock_result.returncode = 1 + mock_result.stderr = b"Error: Image does not exist" + mock_run.return_value = mock_result + + # Act + rc, msg = image_service.set_next_boot(image) + + # Assert + assert rc != 0, "wrong return value" + assert "Error: Image does not exist" in msg, "message should contain 'Error: Image does not exist'" + mock_run.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "set-next-boot", image], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) +