diff --git a/host_modules/image_service.py b/host_modules/image_service.py index e30e1293..c52c6de9 100644 --- a/host_modules/image_service.py +++ b/host_modules/image_service.py @@ -12,6 +12,7 @@ import requests import stat import subprocess +import json from host_modules import host_service import tempfile @@ -133,3 +134,65 @@ def checksum(self, file_path, algorithm): except Exception as e: logger.error("Failed to calculate checksum: {}".format(e)) return errno.EIO, str(e) + + @host_service.method( + host_service.bus_name(MOD_NAME), in_signature="", out_signature="is" + ) + def list_images(self): + """ + List the current, next, and available SONiC images. + + Returns: + A tuple with an error code and a JSON string with keys "current", "next", and "available" or an error message. + """ + logger.info("Listing SONiC images") + + try: + output = subprocess.check_output( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ).decode().strip() + result = self._parse_sonic_installer_list(output) + logger.info("List result: {}".format(result)) + return 0, json.dumps(result) + except subprocess.CalledProcessError as e: + msg = "Failed to list images: command {} failed with return code {} and message {}".format(e.cmd, e.returncode, e.output.decode()) + logger.error(msg) + return e.returncode, msg + + def _parse_sonic_installer_list(self, output): + """ + Parse the output of the sonic-installer list command. + + Args: + output: The output of the sonic-installer list command. + + Returns: + A dictionary with keys "current", "next", and "available" containing the respective images. + """ + current_image = "" + next_image = "" + available_images = [] + + for line in output.split("\n"): + if "current:" in line.lower(): + parts = line.split(":") + if len(parts) > 1: + current_image = parts[1].strip() + elif "next:" in line.lower(): + parts = line.split(":") + if len(parts) > 1: + next_image = parts[1].strip() + elif "available:" in line.lower(): + continue + else: + available_images.append(line.strip()) + + logger.info("Current image: {}".format(current_image)) + logger.info("Next image: {}".format(next_image)) + logger.info("Available images: {}".format(available_images)) + return { + "current": current_image or "", + "next": next_image or "", + "available": available_images or [], + } \ No newline at end of file diff --git a/tests/host_modules/image_service_test.py b/tests/host_modules/image_service_test.py index 22e70e82..bc941778 100644 --- a/tests/host_modules/image_service_test.py +++ b/tests/host_modules/image_service_test.py @@ -4,6 +4,7 @@ import os import stat import pytest +import json from unittest import mock from host_modules.image_service import ImageService @@ -370,3 +371,200 @@ def test_checksum_general_exception( ), "message should contain 'general error'" mock_isfile.assert_called_once_with(file_path) mock_open.assert_called_once_with(file_path, "rb") + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_success(self, mock_check_output, MockInit, MockBusName, MockSystemBus): + """ + Test that the `list_images` method successfully lists the current, next, and available SONiC images. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_output = ( + "Current: current_image\n" + "Next: next_image\n" + "Available:\n" + "image1\n" + "image2\n" + ) + mock_check_output.return_value = mock_output.encode() + + # Act + rc, images_json = image_service.list_images() + images = json.loads(images_json) + + # Assert + assert rc == 0, "wrong return value" + assert images["current"] == "current_image", "current image does not match" + assert images["next"] == "next_image", "next image does not match" + assert images["available"] == ["image1", "image2"], "available images do not match" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_success_lowercase_output(self, mock_check_output, MockInit, MockBusName, MockSystemBus): + """ + Test that the `list_images` method successfully lists the current, next, and available SONiC images + even if the output from sonic-installer is in lowercase. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_output = ( + "current: current_image\n" + "next: next_image\n" + "available:\n" + "image1\n" + "image2\n" + ) + mock_check_output.return_value = mock_output.encode() + + # Act + rc, images_json = image_service.list_images() + images = json.loads(images_json) + + # Assert + assert rc == 0, "wrong return value" + assert images["current"] == "current_image", "current image does not match" + assert images["next"] == "next_image", "next image does not match" + assert images["available"] == ["image1", "image2"], "available images do not match" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) + + @pytest.mark.parametrize( + "mock_output, expected_current, expected_next, expected_available", + [ + ("Current: \nNext: next_image\nAvailable:\nimage1\nimage2\n", "", "next_image", ["image1", "image2"]), + ("Current: current_image\nNext: \nAvailable:\nimage1\nimage2\n", "current_image", "", ["image1", "image2"]), + ("Current: current_image\nNext: next_image\nAvailable:\n", "current_image", "next_image", []), + ], + ) + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_success_empty_image_output( + self, mock_check_output, MockInit, MockBusName, MockSystemBus, mock_output, expected_current, expected_next, expected_available + ): + """ + Test that the `list_images` method successfully lists the current, next, and available SONiC images even if + sonic-installer output empty string for the image name. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_check_output.return_value = mock_output.encode() + + # Act + rc, images_json = image_service.list_images() + images = json.loads(images_json) + + # Assert + assert rc == 0, "wrong return value" + assert images["current"] == expected_current, "current image does not match" + assert images["next"] == expected_next, "next image does not match" + assert images["available"] == expected_available, "available images do not match" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) + + @pytest.mark.parametrize( + "mock_output, expected_current, expected_next, expected_available", + [ + ("Next: next_image\nAvailable:\nimage1\nimage2\n", "", "next_image", ["image1", "image2"]), + ("Current: current_image\nAvailable:\nimage1\nimage2\n", "current_image", "", ["image1", "image2"]), + ("Current: current_image\nNext: next_image\n", "current_image", "next_image", []), + ("Available:\nimage1\nimage2\n", "", "", ["image1", "image2"]), + ("Current: current_image\nNext: next_image\nAvailable:\n", "current_image", "next_image", []), + ], + ) + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_various_missing_lines( + self, mock_check_output, MockInit, MockBusName, MockSystemBus, mock_output, expected_current, expected_next, expected_available + ): + """ + Test that the `list_images` method handles various scenarios where the sonic-installer output is missing lines for current, next, or available images. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_check_output.return_value = mock_output.encode() + + # Act + rc, images_json = image_service.list_images() + images = json.loads(images_json) + + # Assert + assert rc == 0, "wrong return value" + assert images["current"] == expected_current, "current image does not match" + assert images["next"] == expected_next, "next image does not match" + assert images["available"] == expected_available, "available images do not match" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_success_empty_available_images(self, mock_check_output, MockInit, MockBusName, MockSystemBus): + """ + Test that the `list_images` method successfully lists the current, next, and available SONiC images. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_output = ( + "Current: current_image\n" + "Next: next_image\n" + "Available:\n" + ) + mock_check_output.return_value = mock_output.encode() + + # Act + rc, images_json = image_service.list_images() + images = json.loads(images_json) + + # Assert + assert rc == 0, "wrong return value" + assert images["available"] == [], "available images should be empty" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) + + + @mock.patch("dbus.SystemBus") + @mock.patch("dbus.service.BusName") + @mock.patch("dbus.service.Object.__init__") + @mock.patch("subprocess.check_output") + def test_list_images_failed(self, mock_check_output, MockInit, MockBusName, MockSystemBus): + """ + Test that the `list_image` method fails when the subprocess command returns a non-zero exit code. + """ + # Arrange + image_service = ImageService(mod_name="image_service") + mock_check_output.side_effect = subprocess.CalledProcessError( + returncode=1, cmd="sonic-installer list", output=b"Error: command failed" + ) + + # Act + rc, msg = image_service.list_images() + + # Assert + assert rc != 0, "wrong return value" + mock_check_output.assert_called_once_with( + ["/usr/local/bin/sonic-installer", "list"], + stderr=subprocess.STDOUT, + ) +