scripts/xz_auto_unittest.py

# Copyright 2021 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.

"""Unit tests for xz_auto.py."""

import os
from pathlib import Path
import unittest
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.scripts import xz_auto


DIR = Path(__file__).resolve().parent


def FindXzAutoLocation():
  """Figures out where the xz_auto binary is."""
  return DIR / 'xz_auto'


class XzAutoTests(cros_test_lib.MockTempDirTestCase):
  """Various tests for xz_auto."""

  def DisablePixzForCurrentTest(self):
    """Disables the use of pixz for the current test."""
    # This will be cleaned up by cros_test_lib, so no need to addCleanup.
    os.environ[xz_auto.PIXZ_DISABLE_VAR] = '1'

  def testPixzArgParsingSeemsToWork(self):
    """Tests our detection of file names in pixz commandlines."""
    self.assertEqual(
        xz_auto.ParsePixzArgs(['to_compress.txt']),
        ([], 'to_compress.txt', None),
    )
    self.assertEqual(
        xz_auto.ParsePixzArgs(['to_compress.txt', 'compressed.txt']),
        ([], 'to_compress.txt', 'compressed.txt'),
    )
    self.assertEqual(
        xz_auto.ParsePixzArgs(['to_compress.txt', '-c', 'compressed.txt',
                               '-9']),
        (['-c', '-9'], 'to_compress.txt', 'compressed.txt'),
    )
    self.assertEqual(
        xz_auto.ParsePixzArgs(
            ['-t', 'to_compress.txt', '-c', 'compressed.txt', '-p', '2']),
        (['-t', '-c', '-p', '2'], 'to_compress.txt', 'compressed.txt'),
    )
    self.assertEqual(
        xz_auto.ParsePixzArgs(['-tcp2', 'to_compress.txt', 'compressed.txt']),
        (['-t', '-c', '-p', '2'], 'to_compress.txt', 'compressed.txt'),
    )

  @unittest.skipIf(not xz_auto.HasPixz(), 'need pixz for this test')
  @mock.patch.object(xz_auto, 'Execvp')
  def testPixzCommandCreationSelectsPixzIfAvailable(self, execvp_mock):
    """Tests that we actually execute pixz when we intend to."""
    class ExecvpStopError(Exception):
      """Convenient way to halt execution."""

    def execvp_side_effect(argv):
      """Does testing of our execvp calls."""
      self.assertEqual(argv[0], 'pixz')
      raise ExecvpStopError()

    execvp_mock.side_effect = execvp_side_effect
    with self.assertRaises(ExecvpStopError):
      xz_auto.ExecCompressCommand(stdout=False, argv=[])

    with self.assertRaises(ExecvpStopError):
      xz_auto.ExecDecompressCommand(stdout=False, argv=[])

  def _TestFileCompressionImpl(self):
    """Tests that compressing a file with xz_auto WAI."""
    file_contents = b'some random file contents'
    file_location = os.path.join(self.tempdir, 'file.txt')
    osutils.WriteFile(file_location, file_contents, mode='wb')

    xz_auto_script = str(FindXzAutoLocation())
    cros_build_lib.run(
        [
            xz_auto_script,
            file_location,
        ],
        check=True,
    )

    xz_location = file_location + '.xz'
    self.assertExists(xz_location)
    self.assertNotExists(file_location)
    cros_build_lib.run([
        xz_auto_script,
        '--decompress',
        xz_location,
    ])
    self.assertNotExists(xz_location)
    self.assertExists(file_location)
    self.assertEqual(
        osutils.ReadFile(file_location, mode='rb'),
        file_contents,
    )

  def _TestStdoutCompressionImpl(self):
    """Tests that compressing stdstreams with xz_auto WAI."""
    file_contents = b'some random file contents'
    xz_auto_script = str(FindXzAutoLocation())

    run_result = cros_build_lib.run(
        [
            xz_auto_script,
            '-c',
        ],
        capture_output=True,
        input=file_contents,
    )

    compressed_file = run_result.stdout
    self.assertNotEqual(compressed_file, file_contents)

    run_result = cros_build_lib.run(
        [
            xz_auto_script,
            '--decompress',
            '-c',
        ],
        input=compressed_file,
        capture_output=True,
    )
    uncompressed_file = run_result.stdout
    self.assertEqual(file_contents, uncompressed_file)

  def _TestStdoutCompressionFromFileImpl(self):
    """Tests that compression of a file & outputting to stdout works.

    Pixz has some semi-weird behavior here (b/202735786).
    """
    file_contents = b'some random file contents'
    xz_auto_script = str(FindXzAutoLocation())
    file_location = os.path.join(self.tempdir, 'file.txt')
    osutils.WriteFile(file_location, file_contents, mode='wb')

    run_result = cros_build_lib.run(
        [
            xz_auto_script,
            '-c',
            file_location,
        ],
        capture_output=True,
    )

    compressed_file = run_result.stdout
    self.assertExists(file_location)
    self.assertNotEqual(compressed_file, file_contents)

    run_result = cros_build_lib.run(
        [
            xz_auto_script,
            '--decompress',
            '-c',
        ],
        capture_output=True,
        input=compressed_file,
    )
    uncompressed_file = run_result.stdout
    self.assertEqual(file_contents, uncompressed_file)

  @unittest.skipIf(not xz_auto.HasPixz(), 'need pixz for this test')
  def testFileCompressionWithPixzWorks(self):
    """Tests that compressing a file with pixz WAI."""
    self._TestFileCompressionImpl()

  @unittest.skipIf(not xz_auto.HasPixz(), 'need pixz for this test')
  def testStdoutCompressionWithPixzWorks(self):
    """Tests that compressing `stdout` with pixz WAI."""
    self._TestStdoutCompressionImpl()

  @unittest.skipIf(not xz_auto.HasPixz(), 'need pixz for this test')
  def testStdoutCompressionFromFileWithPixzWorks(self):
    """Tests that compressing from a file to stdout with pixz WAI."""
    self._TestStdoutCompressionFromFileImpl()

  def testFileCompressionWithXzWorks(self):
    """Tests that compressing a file with pixz WAI."""
    self.DisablePixzForCurrentTest()
    self._TestFileCompressionImpl()

  def testStdoutCompressionWithXzWorks(self):
    """Tests that compressing `stdout` with pixz WAI."""
    self.DisablePixzForCurrentTest()
    self._TestStdoutCompressionImpl()

  def testStdoutCompressionFromFileWithXzWorks(self):
    """Tests that compressing from a file to stdout WAI."""
    self.DisablePixzForCurrentTest()
    self._TestStdoutCompressionFromFileImpl()