lib/paygen/partition_lib_unittest.py

# -*- coding: utf-8 -*-
# Copyright 2018 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.

"""Test the partition_lib module."""

from __future__ import print_function

import os

from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import osutils

from chromite.lib.paygen import partition_lib

class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
  """Test partition_lib.ExtractRoot"""
  def testTruncate(self):
    """Test truncating on extraction."""
    root = os.path.join(self.tempdir, 'root.bin')
    root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')

    content = '0123456789'
    osutils.WriteFile(root, content)
    self.PatchObject(partition_lib, 'ExtractPartition',
                     return_value=root)
    self.PatchObject(partition_lib, 'Ext2FileSystemSize',
                     return_value=2)

    partition_lib.ExtractRoot(None, root, root_pretruncate)
    self.assertEqual(osutils.ReadFile(root), content[:2])
    self.assertEqual(osutils.ReadFile(root_pretruncate), content)


class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
  """Test partition_lib.Ext2FileSystemSize"""
  def testExt2FileSystemSize(self):
    """Test on simple output."""
    block_size = 4096
    block_count = 123

    self.PatchObject(cros_build_lib, 'RunCommand',
                     return_value=cros_build_lib.CommandResult(output='''
Block size: %d
Other thing: 123456798
Not an integer: cdsad132csda
Block count: %d
''' % (block_size, block_count)))

    size = partition_lib.Ext2FileSystemSize('/dev/null')
    self.assertEqual(size, block_size * block_count)


class ExtractPartitionTest(cros_test_lib.MockTempDirTestCase):
  """Test partition_lib.ExtractPartition"""
  def testExtractPartition(self):
    """Tests extraction on a simple image."""
    part_a_bin = '0123'
    part_b_bin = '4567'
    image_bin = part_a_bin + part_b_bin

    image = os.path.join(self.tempdir, 'image.bin')
    osutils.WriteFile(image, image_bin)
    part_a = os.path.join(self.tempdir, 'a.bin')

    fake_partitions = (
        cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
        cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
    )
    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
                     return_value=fake_partitions)

    partition_lib.ExtractPartition(image, 'PART-A', part_a)
    self.assertEqual(osutils.ReadFile(part_a), part_a_bin)