scripts/cros_sdk_unittest.py

# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Tests for cros_sdk."""

import logging
import os
import re
import subprocess
import sys
import unittest

from chromite.lib import constants
from chromite.lib import cros_build_lib
from chromite.lib import cros_sdk_lib
from chromite.lib import cros_test_lib
from chromite.lib import osutils
from chromite.lib import retry_util
from chromite.lib import sudo
from chromite.scripts import cros_sdk


# This long decorator triggers a false positive in the docstring test.
# https://github.com/PyCQA/pylint/issues/3077
# pylint: disable=bad-docstring-quotes
@unittest.skipIf(cros_build_lib.IsInsideChroot(),
                 'Tests only make sense outside the chroot')
class CrosSdkPrerequisitesTest(cros_test_lib.TempDirTestCase):
  """Tests for required packages on the host machine.

  These are not real unit tests.  If these tests fail, it means your machine is
  missing necessary commands to support image-backed chroots.  Run cros_sdk in
  a terminal to get info about what you need to install.
  """

  def testLvmCommandsPresent(self):
    """Check for commands from the lvm2 package."""
    with sudo.SudoKeepAlive():
      cmd = ['lvs', '--version']
      result = cros_build_lib.run(cmd, check=False)
      self.assertEqual(result.returncode, 0)

  def testThinProvisioningToolsPresent(self):
    """Check for commands from the thin-provisioning-tools package."""
    with sudo.SudoKeepAlive():
      cmd = ['thin_check', '-V']
      result = cros_build_lib.run(cmd, check=False)
      self.assertEqual(result.returncode, 0)

  def testLosetupCommandPresent(self):
    """Check for commands from the mount package."""
    with sudo.SudoKeepAlive():
      cmd = ['losetup', '--help']
      result = cros_build_lib.run(cmd, check=False)
      self.assertEqual(result.returncode, 0)


class CrosSdkUtilsTest(cros_test_lib.MockTempDirTestCase):
  """Tests for misc util funcs."""

  def testGetArchStageTarballs(self):
    """Basic test of GetArchStageTarballs."""
    self.assertCountEqual([
        'https://storage.googleapis.com/chromiumos-sdk/cros-sdk-123.tar.xz',
    ], cros_sdk.GetArchStageTarballs('123'))

  def testFetchRemoteTarballsEmpty(self):
    """Test FetchRemoteTarballs with no results."""
    m = self.PatchObject(retry_util, 'RunCurl')
    with self.assertRaises(ValueError):
      cros_sdk.FetchRemoteTarballs(self.tempdir, [])
    m.return_value = cros_build_lib.CompletedProcess(stdout=b'Foo: bar\n')
    with self.assertRaises(ValueError):
      cros_sdk.FetchRemoteTarballs(self.tempdir, ['gs://x.tar'])

  def testFetchRemoteTarballsSuccess(self):
    """Test FetchRemoteTarballs with a successful download."""
    curl = cros_build_lib.CompletedProcess(stdout=(
        b'HTTP/1.0 200\n'
        b'Foo: bar\n'
        b'Content-Length: 100\n'
    ))
    self.PatchObject(retry_util, 'RunCurl', return_value=curl)
    self.assertEqual(
        os.path.join(self.tempdir, 'tar'),
        cros_sdk.FetchRemoteTarballs(self.tempdir, ['gs://x/tar']))


class CrosSdkParserCommandLineTest(cros_test_lib.MockTestCase):
  """Tests involving the CLI."""

  # pylint: disable=protected-access

  # A typical sys.argv[0] that cros_sdk sees.
  ARGV0 = '/home/chronos/chromiumos/chromite/bin/cros_sdk'

  def setUp(self):
    self.parser, _ = cros_sdk._CreateParser('1', '2')

  def testSudoCommand(self):
    """Verify basic sudo command building works."""
    # Stabilize the env for testing.
    for v in constants.CHROOT_ENVIRONMENT_ALLOWLIST + constants.ENV_PASSTHRU:
      os.environ[v] = 'value'
    os.environ['PATH'] = 'path'

    cmd = cros_sdk._SudoCommand()
    assert cmd[0] == 'sudo'
    assert 'CHROMEOS_SUDO_PATH=path' in cmd
    rlimits = [x for x in cmd if x.startswith('CHROMEOS_SUDO_RLIMITS=')]
    assert len(rlimits) == 1

    # Spot check some pass thru vars.
    assert 'GIT_AUTHOR_EMAIL=value' in cmd
    assert 'https_proxy=value' in cmd

    # Make sure we only pass vars after `sudo`.
    for i in range(1, len(cmd)):
      assert '=' in cmd[i]
      v = cmd[i].split('=', 1)[0]
      assert re.match(r'^[A-Za-z0-9_]+$', v) is not None

  def testReexecCommand(self):
    """Verify reexec command line building."""
    # Stub sudo logic since we tested it above already.
    self.PatchObject(cros_sdk, '_SudoCommand', return_value=['sudo'])
    opts = self.parser.parse_args([])
    new_cmd = cros_sdk._BuildReExecCommand([self.ARGV0], opts)
    assert new_cmd == ['sudo', '--', sys.executable, self.ARGV0]

  def testReexecCommandStrace(self):
    """Verify reexec command line building w/strace."""
    # Stub sudo logic since we tested it above already.
    self.PatchObject(cros_sdk, '_SudoCommand', return_value=['sudo'])

    # Strace args passed, but not enabled.
    opts = self.parser.parse_args(['--strace-arguments=-s4096 -v'])
    new_cmd = cros_sdk._BuildReExecCommand([self.ARGV0], opts)
    assert new_cmd == ['sudo', '--', sys.executable, self.ARGV0]

    # Strace enabled.
    opts = self.parser.parse_args(['--strace'])
    new_cmd = cros_sdk._BuildReExecCommand([self.ARGV0], opts)
    assert new_cmd == ['sudo', '--', 'strace', '--', sys.executable, self.ARGV0]

    # Strace enabled w/args.
    opts = self.parser.parse_args(['--strace', '--strace-arguments=-s4096 -v'])
    new_cmd = cros_sdk._BuildReExecCommand([self.ARGV0], opts)
    assert new_cmd == [
        'sudo', '--', 'strace', '-s4096', '-v', '--', sys.executable,
        self.ARGV0,
    ]


@unittest.skipIf(cros_build_lib.IsInsideChroot(),
                 'Tests only make sense outside the chroot')
@unittest.skip(
    'These are triggering hangs inside lvm when run through run_tests. '
    'https://crbug.com/764335')
class CrosSdkSnapshotTest(cros_test_lib.TempDirTestCase):
  """Tests for the snapshot functionality in cros_sdk."""

  def setUp(self):
    with sudo.SudoKeepAlive():
      # Create just enough of a chroot to fool cros_sdk into accepting it.
      self.chroot = os.path.join(self.tempdir, 'chroot')
      cros_sdk_lib.MountChroot(self.chroot, create=True)
      logging.debug('Chroot mounted on %s', self.chroot)

      chroot_etc = os.path.join(self.chroot, 'etc')
      osutils.SafeMakedirsNonRoot(chroot_etc)

      self.chroot_version_file = os.path.join(chroot_etc, 'cros_chroot_version')
      osutils.Touch(self.chroot_version_file, makedirs=True)

  def tearDown(self):
    with sudo.SudoKeepAlive():
      cros_sdk_lib.CleanupChrootMount(self.chroot, delete=True)

  def _crosSdk(self, args):
    cmd = ['cros_sdk', '--chroot', self.chroot]
    cmd.extend(args)

    try:
      result = cros_build_lib.run(
          cmd, print_cmd=False, stdout=True, check=False,
          stderr=subprocess.STDOUT)
    except cros_build_lib.RunCommandError as e:
      raise SystemExit('Running %r failed!: %s' % (cmd, e))

    return result.returncode, result.stdout

  def testSnapshotsRequireImage(self):
    code, output = self._crosSdk(['--snapshot-list', '--nouse-image'])
    self.assertNotEqual(code, 0)
    self.assertIn('Snapshot operations are not compatible with', output)

    code, output = self._crosSdk(['--snapshot-delete', 'test', '--nouse-image'])
    self.assertNotEqual(code, 0)
    self.assertIn('Snapshot operations are not compatible with', output)

    code, output = self._crosSdk(['--snapshot-create', 'test', '--nouse-image'])
    self.assertNotEqual(code, 0)
    self.assertIn('Snapshot operations are not compatible with', output)

    code, output = self._crosSdk(['--snapshot-restore', 'test',
                                  '--nouse-image'])
    self.assertNotEqual(code, 0)
    self.assertIn('Snapshot operations are not compatible with', output)

  def testSnapshotWithDeleteChroot(self):
    code, output = self._crosSdk(['--delete', '--snapshot-list'])
    self.assertNotEqual(code, 0)
    self.assertIn('Trying to enter or snapshot the chroot', output)

    code, output = self._crosSdk(['--delete', '--snapshot-delete', 'test'])
    self.assertNotEqual(code, 0)
    self.assertIn('Trying to enter or snapshot the chroot', output)

    code, output = self._crosSdk(['--delete', '--snapshot-create', 'test'])
    self.assertNotEqual(code, 0)
    self.assertIn('Trying to enter or snapshot the chroot', output)

    code, output = self._crosSdk(['--delete', '--snapshot-restore', 'test'])
    self.assertNotEqual(code, 0)
    self.assertIn('Trying to enter or snapshot the chroot', output)

  def testEmptySnapshotList(self):
    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output, '')

  def testOneSnapshot(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output.strip(), 'test')

  def testMultipleSnapshots(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    code, _ = self._crosSdk(['--snapshot-create', 'test2'])
    self.assertEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertSetEqual(set(output.strip().split('\n')), {'test', 'test2'})

  def testCantCreateSameSnapshotTwice(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertNotEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output.strip(), 'test')

  def testCreateSnapshotMountsAsNeeded(self):
    with sudo.SudoKeepAlive():
      cros_sdk_lib.CleanupChrootMount(self.chroot)

    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    self.assertExists(self.chroot_version_file)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output.strip(), 'test')

  def testDeleteGoodSnapshot(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    code, _ = self._crosSdk(['--snapshot-create', 'test2'])
    self.assertEqual(code, 0)

    code, _ = self._crosSdk(['--snapshot-delete', 'test'])
    self.assertEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output.strip(), 'test2')

  def testDeleteMissingSnapshot(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    code, _ = self._crosSdk(['--snapshot-create', 'test2'])
    self.assertEqual(code, 0)

    code, _ = self._crosSdk(['--snapshot-delete', 'test3'])
    self.assertEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertSetEqual(set(output.strip().split('\n')), {'test', 'test2'})

  def testDeleteAndCreateSnapshot(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)
    code, _ = self._crosSdk(['--snapshot-create', 'test2'])
    self.assertEqual(code, 0)

    code, _ = self._crosSdk(['--snapshot-create', 'test',
                             '--snapshot-delete', 'test'])
    self.assertEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertSetEqual(set(output.strip().split('\n')), {'test', 'test2'})

  def testRestoreSnapshot(self):
    with sudo.SudoKeepAlive():
      test_file = os.path.join(self.chroot, 'etc', 'test_file')
      osutils.Touch(test_file)

      code, _ = self._crosSdk(['--snapshot-create', 'test'])
      self.assertEqual(code, 0)

      osutils.SafeUnlink(test_file)

      code, _ = self._crosSdk(['--snapshot-restore', 'test'])
      self.assertEqual(code, 0)
      self.assertTrue(cros_sdk_lib.MountChroot(self.chroot, create=False))
      self.assertExists(test_file)

      code, output = self._crosSdk(['--snapshot-list'])
      self.assertEqual(code, 0)
      self.assertEqual(output, '')

  def testRestoreAndCreateSnapshot(self):
    with sudo.SudoKeepAlive():
      test_file = os.path.join(self.chroot, 'etc', 'test_file')
      osutils.Touch(test_file)

      code, _ = self._crosSdk(['--snapshot-create', 'test'])
      self.assertEqual(code, 0)

      osutils.SafeUnlink(test_file)

      code, _ = self._crosSdk(['--snapshot-restore', 'test',
                               '--snapshot-create', 'test'])
      self.assertEqual(code, 0)
      self.assertTrue(cros_sdk_lib.MountChroot(self.chroot, create=False))
      self.assertExists(test_file)

      code, output = self._crosSdk(['--snapshot-list'])
      self.assertEqual(code, 0)
      self.assertEqual(output.strip(), 'test')

  def testDeleteCantRestoreSameSnapshot(self):
    code, _ = self._crosSdk(['--snapshot-create', 'test'])
    self.assertEqual(code, 0)

    code, _ = self._crosSdk(['--snapshot-delete', 'test',
                             '--snapshot-restore', 'test'])
    self.assertNotEqual(code, 0)

    code, output = self._crosSdk(['--snapshot-list'])
    self.assertEqual(code, 0)
    self.assertEqual(output.strip(), 'test')

  def testCantRestoreInvalidSnapshot(self):
    with sudo.SudoKeepAlive():
      test_file = os.path.join(self.chroot, 'etc', 'test_file')
      osutils.Touch(test_file)

    code, _ = self._crosSdk(['--snapshot-restore', 'test'])
    self.assertNotEqual(code, 0)
    # Failed restore leaves the existing snapshot in place.
    self.assertExists(test_file)

  def testRestoreSnapshotMountsAsNeeded(self):
    with sudo.SudoKeepAlive():
      test_file = os.path.join(self.chroot, 'etc', 'test_file')
      osutils.Touch(test_file)

      code, _ = self._crosSdk(['--snapshot-create', 'test'])
      self.assertEqual(code, 0)

      osutils.SafeUnlink(test_file)

      cros_sdk_lib.CleanupChrootMount(self.chroot)

      code, _ = self._crosSdk(['--snapshot-restore', 'test'])
      self.assertEqual(code, 0)
      self.assertExists(test_file)

      code, output = self._crosSdk(['--snapshot-list'])
      self.assertEqual(code, 0)
      self.assertEqual(output, '')