updated unit test for download

tests/host_modules/image_service_test.py

@@ -1,5 +1,6 @@
 import sys
 import os
+import stat
 import pytest
 from unittest import mock
 from host_modules.image_service import ImageService
@@ -9,184 +10,119 @@ class TestImageService(object):
     @mock.patch("dbus.SystemBus")
     @mock.patch("dbus.service.BusName")
     @mock.patch("dbus.service.Object.__init__")
-    def test_download_success(self, MockInit, MockBusName, MockSystemBus):
+    @mock.patch("os.path.isdir")
+    @mock.patch("os.stat")
+    @mock.patch("requests.get")
+    def test_download_success(
+        self, mock_get, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus
+    ):
         """
-        Test that the `download_sonic_image` method runs the correct curl command
-        when the directory path exists.
+        Test that the `download` method successfully downloads an image when the directory exists and is writable.
         """
-        with (
-            mock.patch("os.path.exists", return_value=True) as mock_exists,
-            mock.patch("subprocess.run") as mock_run,
-        ):
-            # Arrange: Set up an Installer instance and define the target path and URL
-            image_service = ImageService(mod_name="image_service")
-            target_path = "/path/to/sonic_image.img"
-            image_url = "https://example.com/sonic_image.img"
-            run_ret = mock.Mock()
-            attrs = {"returncode": 0, "stderr": b""}
-            run_ret.configure_mock(**attrs)
-            mock_run.return_value = run_ret
-
-            # Act: Call the method to download the image
-            rc, msg = image_service.download(image_url, target_path)
-
-            # Assert: Verify that os.path.exists was called to check directory existence
-            mock_exists.assert_called_once_with(os.path.dirname(target_path))
-            assert rc == 0, "wrong return value"
-            assert msg == "", "non-empty return message"
-
-            # Assert: Verify that subprocess.run was called with the correct curl command
-            mock_run.assert_called_once_with(
-                ["/usr/bin/curl", "-Lo", target_path, image_url],
-                stdout=mock.ANY,
-                stderr=mock.ANY,
-            )
+        # 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_response.iter_content = lambda chunk_size: iter([b"data"])
+        mock_get.return_value = mock_response
+
+        # Act
+        rc, msg = image_service.download(image_url, save_as)
+
+        # Assert
+        assert rc == 0, "wrong return value"
+        assert (
+            "download" in msg.lower() and "successful" in msg.lower()
+        ), "message should contains 'download' and 'successful'"
+        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__")
-    def test_download_mkdir(self, MockInit, MockBusName, MockSystemBus):
+    @mock.patch("os.path.isdir")
+    def test_download_fail_no_dir(
+        self, mock_isdir, MockInit, MockBusName, MockSystemBus
+    ):
         """
-        Test that the `download_sonic_image` method runs the correct curl command
-        when the directory path exists.
+        Test that the `download` method fails when the directory does not exist.
         """
-        with (
-            mock.patch("os.path.exists", return_value=False) as mock_exists,
-            mock.patch("os.makedirs") as mock_mkdirs,
-            mock.patch("subprocess.run", return_value=0) as mock_run,
-        ):
-            # Arrange: Set up an Installer instance and define the target path and URL
-            image_service = ImageService(mod_name="image_service")
-            target_path = "/path/to/sonic_image.img"
-            image_url = "https://example.com/sonic_image.img"
-            run_ret = mock.Mock()
-            attrs = {"returncode": 0, "stderr": b""}
-            run_ret.configure_mock(**attrs)
-            mock_run.return_value = run_ret
+        # Arrange
+        image_service = ImageService(mod_name="image_service")
+        image_url = "http://example.com/sonic_image.img"
+        save_as = "/nonexistent_dir/sonic_image.img"
+        mock_isdir.return_value = False
 
-            # Act: Call the method to download the image
-            rc, msg = image_service.download(image_url, target_path)
-            assert rc == 0, "wrong return value"
-            assert msg == "", "non-empty return message"
+        # Act
+        rc, msg = image_service.download(image_url, save_as)
 
-            # Assert: Verify that os.path.exists was called to check directory existence
-            mock_exists.assert_called_once_with(os.path.dirname(target_path))
-            mock_mkdirs.assert_called_once_with(os.path.dirname(target_path))
-            # Assert: Verify that subprocess.run was called with the correct curl command
-            mock_run.assert_called_once_with(
-                ["/usr/bin/curl", "-Lo", target_path, image_url],
-                stdout=mock.ANY,
-                stderr=mock.ANY,
-            )
+        # Assert
+        assert rc != 0, "wrong return value"
+        assert (
+            "not" in msg.lower() and "exist" in msg.lower()
+        ), "message should contains 'not' and 'exist'"
 
     @mock.patch("dbus.SystemBus")
     @mock.patch("dbus.service.BusName")
     @mock.patch("dbus.service.Object.__init__")
-    def test_download_download_fail(self, MockInit, MockBusName, MockSystemBus):
+    @mock.patch("os.path.isdir")
+    @mock.patch("os.stat")
+    def test_download_fail_missing_other_write(
+        self, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus
+    ):
         """
-        Test that the `download_sonic_image` method runs the correct curl command
-        when the directory path exists.
+        Test that the `download` method fails when the directory is not writable by others.
         """
-        with (
-            mock.patch("os.path.exists", return_value=False) as mock_exists,
-            mock.patch("os.makedirs") as mock_mkdirs,
-            mock.patch("subprocess.run", return_value=0) as mock_run,
-        ):
-            # Arrange: Set up an Installer instance and define the target path and URL
-            image_service = ImageService(mod_name="image_service")
-            target_path = "/path/to/sonic_image.img"
-            image_url = "https://example.com/sonic_image.img"
-            run_ret = mock.Mock()
-            # Download failed.
-            attrs = {"returncode": 1, "stderr": b"Error: Download failed\nHello World!"}
-            run_ret.configure_mock(**attrs)
-            mock_run.return_value = run_ret
-
-            # Act: Call the method to download the image
-            rc, msg = image_service.download(image_url, target_path)
-            assert rc != 0, "wrong return value"
-            assert "Error" in msg, "return message without error"
-
-            # Assert: Verify that os.path.exists was called to check directory existence
-            mock_exists.assert_called_once_with(os.path.dirname(target_path))
-            mock_mkdirs.assert_called_once_with(os.path.dirname(target_path))
-            # Assert: Verify that subprocess.run was called with the correct curl command
-            mock_run.assert_called_once_with(
-                ["/usr/bin/curl", "-Lo", target_path, image_url],
-                stdout=mock.ANY,
-                stderr=mock.ANY,
-            )
+        # 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
+        )  # Missing write permission for others
+
+        # Act
+        rc, msg = image_service.download(image_url, save_as)
+
+        # Assert
+        assert rc != 0, "wrong return value"
+        assert (
+            "permission" in msg.lower() or "writable" in msg.lower()
+        ), "message should contain 'permission' or 'writable'"
 
     @mock.patch("dbus.SystemBus")
     @mock.patch("dbus.service.BusName")
     @mock.patch("dbus.service.Object.__init__")
-    def test_install_failed(self, MockInit, MockBusName, MockSystemBus):
+    @mock.patch("os.path.isdir")
+    @mock.patch("os.stat")
+    @mock.patch("requests.get")
+    def test_download_failed_not_found(
+        self, mock_get, mock_stat, mock_isdir, MockInit, MockBusName, MockSystemBus
+    ):
         """
-        Test that the `download_sonic_image` method runs the correct curl command
-        when the directory path exists.
+        Test that the `download` method fails when the image URL is not found (404 error).
         """
-        with mock.patch("subprocess.run") as mock_run:
-            # Arrange: Set up an Installer instance and define the target path and URL
-            image_service = ImageService(mod_name="image_service")
-            target_path = "/path/to/sonic_image.img"
-            run_ret = mock.Mock()
-            attrs = {"returncode": 1, "stderr": b"Error: Install failed\nHello World!"}
-            run_ret.configure_mock(**attrs)
-            mock_run.return_value = run_ret
-
-            # Act: Call the method to download the image
-            rc, msg = image_service.install(target_path)
-
-            # Assert: Verify that os.path.exists was called to check directory existence
-            assert rc != 0, "wrong return value"
-            assert "Error" in msg, "wrong return message"
-
-            # Assert: Verify that subprocess.run was called with the correct curl command
-            mock_run.assert_called_once_with(
-                [
-                    "sudo",
-                    "/usr/local/bin/sonic-installer",
-                    "install",
-                    "-y",
-                    target_path,
-                ],
-                stdout=mock.ANY,
-                stderr=mock.ANY,
-            )
-
-    @mock.patch("dbus.SystemBus")
-    @mock.patch("dbus.service.BusName")
-    @mock.patch("dbus.service.Object.__init__")
-    def test_install_success(self, MockInit, MockBusName, MockSystemBus):
-        """
-        Test that the `download_sonic_image` method runs the correct curl command
-        when the directory path exists.
-        """
-        with mock.patch("subprocess.run") as mock_run:
-            # Arrange: Set up an Installer instance and define the target path and URL
-            image_service = ImageService(mod_name="image_service")
-            target_path = "/path/to/sonic_image.img"
-            run_ret = mock.Mock()
-            attrs = {"returncode": 0, "stderr": b""}
-            run_ret.configure_mock(**attrs)
-            mock_run.return_value = run_ret
-
-            # Act: Call the method to download the image
-            rc, msg = image_service.install(target_path)
-
-            # Assert: Verify that os.path.exists was called to check directory existence
-            assert rc == 0, "wrong return value"
-            assert msg == "", "non-empty return message"
-
-            # Assert: Verify that subprocess.run was called with the correct curl command
-            mock_run.assert_called_once_with(
-                [
-                    "sudo",
-                    "/usr/local/bin/sonic-installer",
-                    "install",
-                    "-y",
-                    target_path,
-                ],
-                stdout=mock.ANY,
-                stderr=mock.ANY,
-            )
+        # Arrange
+        image_service = ImageService(mod_name="image_service")
+        image_url = "http://example.com/nonexistent_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.status_code = 404
+        mock_get.return_value = mock_response
+
+        # Act
+        rc, msg = image_service.download(image_url, save_as)
+
+        # Assert
+        assert rc != 0, "wrong return value"
+        assert (
+            "404" in msg and "error" in msg.lower()
+        ), "message should contain '404' and 'error'"
+        mock_get.assert_called_once_with(image_url, stream=True)