Add unit tests so we have 100% coverage.

host_modules/image_service.py

@@ -65,8 +65,8 @@ def download(self, image_url, save_as):
                     f.write(chunk)
             os.rename(TMP_IMAGE_FILE, save_as)
             return 0, "Download successful"
-        except requests.exceptions.RequestException as e:
-            logger.error("Failed to download image: {}".format(e))
+        except Exception as e:
+            logger.error("Failed to write downloaded image to disk: {}".format(e))
             return errno.EIO, str(e)
 
     @host_service.method(

tests/host_modules/image_service_test.py

@@ -150,6 +150,40 @@ def test_download_failed_not_found(
         ), "message should contain '404' and 'error'"
         mock_get.assert_called_once_with(image_url, stream=True)
 
+    @mock.patch("dbus.SystemBus")
+    @mock.patch("dbus.service.BusName")
+    @mock.patch("dbus.service.Object.__init__")
+    @mock.patch("os.path.isdir")
+    @mock.patch("os.stat")
+    @mock.patch("requests.get")
+    @mock.patch("builtins.open", new_callable=mock.mock_open)
+    def test_download_fail_write_io_exception(
+        self, mock_open, mock_get, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus
+    ):
+        """
+        Test that the `download` method fails when there is an IOError while writing the file.
+        """
+        # Arrange
+        image_service = ImageService(mod_name="image_service")
+        image_url = "http://example.com/sonic_image.img"
+        save_as = "/tmp/sonic_image.img"
+        mock_isdir.return_value = True
+        mock_stat.return_value.st_mode = stat.S_IWUSR | stat.S_IWGRP | stat.S_IWOTH
+        mock_response = mock.Mock()
+        mock_response.iter_content = lambda chunk_size: [b"data"]
+        mock_response.status_code = 200
+        mock_get.return_value = mock_response
+        mock_open.side_effect = IOError("Disk write error")
+
+        # Act
+        rc, msg = image_service.download(image_url, save_as)
+
+        # Assert
+        assert rc != 0, "wrong return value"
+        assert "disk write error" in msg.lower(), "message should contain 'disk write error'"
+        mock_get.assert_called_once_with(image_url, stream=True)
+        mock_open.assert_called_once_with("/tmp/tmp-sonic_image.bin", "wb")
+
     @mock.patch("dbus.SystemBus")
     @mock.patch("dbus.service.BusName")
     @mock.patch("dbus.service.Object.__init__")
@@ -269,3 +303,57 @@ def test_checksum_no_such_file(
         assert rc != 0, "wrong return value"
         assert "not exist" in msg.lower(), "message should contain 'not exist'"
         mock_isfile.assert_called_once_with(file_path)
+
+    @mock.patch("dbus.SystemBus")
+    @mock.patch("dbus.service.BusName")
+    @mock.patch("dbus.service.Object.__init__")
+    @mock.patch("os.path.isfile")
+    def test_checksum_unsupported_algorithm(
+        self, mock_isfile, MockInit, MockBusName, MockSystemBus
+    ):
+        """
+        Test that the `checksum` method fails when an unsupported algorithm is provided.
+        """
+        # Arrange
+        image_service = ImageService(mod_name="image_service")
+        file_path = "/tmp/test_file.img"
+        algorithm = "unsupported_algo"
+        mock_isfile.return_value = True
+
+        # Act
+        rc, msg = image_service.checksum(file_path, algorithm)
+
+        # Assert
+        assert rc != 0, "wrong return value"
+        assert "unsupported algorithm" in msg.lower(), "message should contain 'unsupported algorithm'"
+        mock_isfile.assert_called_once_with(file_path)
+
+    @mock.patch("dbus.SystemBus")
+    @mock.patch("dbus.service.BusName")
+    @mock.patch("dbus.service.Object.__init__")
+    @mock.patch("os.path.isfile")
+    @mock.patch("builtins.open", new_callable=mock.mock_open, read_data=b"test data")
+    def test_checksum_general_exception(
+        self, mock_open, mock_isfile, MockInit, MockBusName, MockSystemBus
+    ):
+        """
+        Test that the `checksum` method handles general exceptions during file reading.
+        """
+        # Arrange
+        image_service = ImageService(mod_name="image_service")
+        file_path = "/tmp/test_file.img"
+        algorithm = "sha256"
+        mock_isfile.return_value = True
+
+        with mock.patch.object(hashlib, algorithm) as mock_hash_func:
+            mock_hash_instance = mock_hash_func.return_value
+            mock_hash_instance.update.side_effect = Exception("General error")
+
+            # Act
+            rc, msg = image_service.checksum(file_path, algorithm)
+
+            # Assert
+            assert rc != 0, "wrong return value"
+            assert "general error" in msg.lower(), "message should contain 'general error'"
+            mock_isfile.assert_called_once_with(file_path)
+            mock_open.assert_called_once_with(file_path, "rb")