lib/stateful_updater_unittest.py

# Copyright 2020 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Unit tests for StatefulUpdater."""

import os
import shutil
from unittest import mock

from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import osutils
from chromite.lib import partial_mock
from chromite.lib import remote_access
from chromite.lib import stateful_updater


# pylint: disable=protected-access


class ChromiumOSDeviceMock(partial_mock.PartialMock):
    """Provides a context where RemoteDevice run commands are run locally."""

    TARGET = "chromite.lib.remote_access.ChromiumOSDevice"
    ATTRS = ("run", "CopyToDevice")

    # pylint: disable=unused-argument
    def run(self, inst, *args, **kwargs):
        """Partial mock for ChromiumOSDevice.run."""
        return cros_build_lib.run(*args, **kwargs)

    def CopyToDevice(self, inst, src, dst, mode, **kwargs):
        """Partial mock for ChromiumOSDevice.CopyToDevice."""
        return shutil.copy2(src, dst)


class StatefulUpdaterTest(cros_test_lib.MockTempDirTestCase):
    """A class for testing StatefulUpdater."""

    def setUp(self) -> None:
        """Sets up the objects needed for testing."""
        self._CreateStatefulUpdate()
        self._stateful_dir = os.path.join(self.tempdir, "target")
        osutils.SafeMakedirs(self._stateful_dir)

        self.StartPatcher(ChromiumOSDeviceMock())

    def _CreateStatefulUpdate(self) -> None:
        """Creates a stateful update tar file so we can test it."""
        self._payload = os.path.join(self.tempdir, "stateful.tgz")

        tmp_stateful = os.path.join(self.tempdir, "temp-stateful")
        stateful_dirs = ("var_new", "dev_image_new")
        for d in stateful_dirs:
            osutils.SafeMakedirs(os.path.join(tmp_stateful, d))

        cros_build_lib.CreateTarball(
            self._payload,
            tmp_stateful,
            compression=cros_build_lib.CompressionType.GZIP,
            inputs=stateful_dirs,
        )
        self.assertExists(self._payload)

    def testTargetStatefulUpdateFileDoNotExist(self) -> None:
        """Verify we raise error if the target stateful file doesn't exist."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(device)
            with self.assertRaises(stateful_updater.Error):
                updater.Update("/foo/path")

    def testUpdateStandard(self) -> None:
        """Tests Update function with default arguments."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(
                device, stateful_dir=self._stateful_dir
            )

            updater.Update(self._payload)
            self.assertEqual(
                osutils.ReadFile(
                    os.path.join(self._stateful_dir, updater._UPDATE_TYPE_FILE)
                ),
                "",
            )

    def testUpdateFileDescriptor(self) -> None:
        """Tests Update function with file descriptor input."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(
                device, stateful_dir=self._stateful_dir
            )

            r, w = os.pipe()
            with os.fdopen(w, "wb") as fp:
                fp.write(osutils.ReadFile(self._payload, "rb"))

            updater.Update(r, is_payload_on_device=False)
            os.close(r)
            self.assertEqual(
                osutils.ReadFile(
                    os.path.join(self._stateful_dir, updater._UPDATE_TYPE_FILE)
                ),
                "",
            )

    def testUpdateClobber(self) -> None:
        """Tests Update function with default arguments."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(
                device, stateful_dir=self._stateful_dir
            )

            updater.Update(
                self._payload, update_type=updater.UPDATE_TYPE_CLOBBER
            )
            self.assertEqual(
                osutils.ReadFile(
                    os.path.join(self._stateful_dir, updater._UPDATE_TYPE_FILE)
                ),
                updater.UPDATE_TYPE_CLOBBER,
            )

    def testInvalidUpdateType(self) -> None:
        """Tests we raise error on invalid given update type."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(device)
            with self.assertRaises(stateful_updater.Error):
                updater._MarkUpdateType("foo")

    @mock.patch.object(
        remote_access.RemoteDevice, "IfPathExists", return_value=False
    )
    def testUpdateFailed(self, _) -> None:
        """Tests Update function fails to untar the payload."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(
                device, stateful_dir=self._stateful_dir
            )

            with self.assertRaises(stateful_updater.Error):
                updater.Update(self._payload)

    def testReset(self) -> None:
        """Tests Reset function."""
        with remote_access.ChromiumOSDeviceHandler(
            remote_access.TEST_IP
        ) as device:
            updater = stateful_updater.StatefulUpdater(
                device, stateful_dir=self._stateful_dir
            )
            expected_files = (
                os.path.join(self.tempdir, "target", f)
                for f in (
                    updater._VAR_DIR,
                    updater._DEV_IMAGE_DIR,
                    updater._UPDATE_TYPE_FILE,
                )
            )

            updater.Update(self._payload)

            for f in expected_files:
                self.assertExists(f)

            updater.Reset()

            # These files/directories should be deleted after the reset.
            for f in expected_files:
                self.assertNotExists(f)