#!/usr/bin/python
# Copyright (c) 2012 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.

"""Unittests for the osutils.py module (imagine that!)."""

import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.abspath(__file__)))))

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

# TODO(build): Finish test wrapper (http://crosbug.com/37517).
# Until then, this has to be after the chromite imports.
import mock

class TestOsutils(cros_test_lib.TempDirTestCase):
  """General unittests for the osutils module."""

  def testReadWriteFile(self):
    """Verify we can write data to a file, and then read it back."""
    filename = os.path.join(self.tempdir, 'foo')
    data = 'alsdkfjasldkfjaskdlfjasdf'
    self.assertEqual(osutils.WriteFile(filename, data), None)
    self.assertEqual(osutils.ReadFile(filename), data)

  def testSafeUnlink(self):
    """Test unlinking files work (existing or not)."""
    def f(dirname, sudo=False):
      dirname = os.path.join(self.tempdir, dirname)
      path = os.path.join(dirname, 'foon')
      os.makedirs(dirname)
      open(path, 'w').close()
      self.assertTrue(os.path.exists(path))
      if sudo:
        cros_build_lib.SudoRunCommand(
            ['chown', 'root:root', '-R', '--', dirname], print_cmd=False)
        self.assertRaises(EnvironmentError, os.unlink, path)
      self.assertTrue(osutils.SafeUnlink(path, sudo=sudo))
      self.assertFalse(os.path.exists(path))
      self.assertFalse(osutils.SafeUnlink(path))
      self.assertFalse(os.path.exists(path))

    f("nonsudo", False)
    f("sudo", True)

  def testSafeMakedirs(self):
    """Test creating directory trees work (existing or not)."""
    path = os.path.join(self.tempdir, 'a', 'b', 'c', 'd', 'e')
    self.assertTrue(osutils.SafeMakedirs(path))
    self.assertTrue(os.path.exists(path))
    self.assertFalse(osutils.SafeMakedirs(path))
    self.assertTrue(os.path.exists(path))

  def testSafeMakedirs_error(self):
    """Check error paths."""
    self.assertRaises(OSError, osutils.SafeMakedirs, '/foo/bar/cow/moo/wee')
    self.assertRaises(OSError, osutils.SafeMakedirs, '')

  def testSafeMakedirsSudo(self):
    """Test creating directory trees work as root (existing or not)."""
    path = os.path.join(self.tempdir, 'a', 'b', 'c', 'd', 'e')
    self.assertTrue(osutils.SafeMakedirs(path, sudo=True))
    self.assertTrue(os.path.exists(path))
    self.assertFalse(osutils.SafeMakedirs(path, sudo=True))
    self.assertTrue(os.path.exists(path))
    self.assertEqual(os.stat(path).st_uid, 0)
    # Have to manually clean up as a non-root `rm -rf` will fail.
    cros_build_lib.SudoRunCommand(['rm', '-rf', self.tempdir], print_cmd=False)

  def testRmDir(self):
    """Test that removing dirs work."""
    path = os.path.join(self.tempdir, 'a', 'b', 'c', 'd', 'e')

    self.assertRaises(EnvironmentError, osutils.RmDir, path)
    osutils.SafeMakedirs(path)
    osutils.RmDir(path)
    osutils.RmDir(path, ignore_missing=True)
    self.assertRaises(EnvironmentError, osutils.RmDir, path)

    osutils.SafeMakedirs(path)
    osutils.RmDir(path)
    self.assertFalse(os.path.exists(path))

  def testRmDirSudo(self):
    """Test that removing dirs via sudo works."""
    subpath = os.path.join(self.tempdir, 'a')
    path = os.path.join(subpath, 'b', 'c', 'd', 'e')
    self.assertTrue(osutils.SafeMakedirs(path, sudo=True))
    self.assertRaises(OSError, osutils.RmDir, path)
    osutils.RmDir(subpath, sudo=True)
    self.assertRaises(cros_build_lib.RunCommandError,
                      osutils.RmDir, subpath, sudo=True)

  def testTouchFile(self):
    """Test that we can touch files."""
    path = os.path.join(self.tempdir, 'touchit')
    self.assertFalse(os.path.exists(path))
    osutils.Touch(path)
    self.assertTrue(os.path.exists(path))
    self.assertEqual(os.path.getsize(path), 0)

  def testTouchFileSubDir(self):
    """Test that we can touch files in non-existent subdirs."""
    path = os.path.join(self.tempdir, 'a', 'b', 'c', 'touchit')
    self.assertFalse(os.path.exists(os.path.dirname(path)))
    osutils.Touch(path, makedirs=True)
    self.assertTrue(os.path.exists(path))
    self.assertEqual(os.path.getsize(path), 0)

  def testFindDepotTools(self):
    """Verify FindDepotTools() works in various cases."""
    depot_tools_path = self.tempdir
    gclient_path = os.path.join(self.tempdir, 'gclient.py')
    gitcl_path = os.path.join(self.tempdir, 'git_cl.py')

    path = ':'.join(['/usr/local/bin', depot_tools_path, '/usr/bin'])
    os.environ['PATH'] = path
    # No binaries found.
    self.assertEquals(osutils.FindDepotTools(), None)

    osutils.Touch(gclient_path)
    # No git_cl.py in path.
    self.assertEquals(osutils.FindDepotTools(), None)

    osutils.Touch(gitcl_path)
    self.assertEquals(osutils.FindDepotTools(), depot_tools_path)

    # Check trailing '/' in depot tools location in PATH is handled
    # correctly.
    path = ':'.join(['/usr/local/bin', depot_tools_path + '/', '/usr/bin'])
    os.environ['PATH'] = path
    self.assertEquals(osutils.FindDepotTools(), depot_tools_path)


class TempDirTests(cros_test_lib.TestCase):
  """Unittests of osutils.TempDir.

  Unlike other test classes in this file, TempDirTestCase isn't used as a base
  class, because that is the functionality under test.
  """
  PREFIX = 'chromite.test.osutils.TempDirTests'

  class HelperException(Exception):
    """Exception for tests to raise to test exception handling."""

  class HelperExceptionInner(Exception):
    """Exception for tests to raise to test exception handling."""

  def testBasicSuccessEmpty(self):
    """Test we create and cleanup an empty tempdir."""
    with osutils.TempDir(prefix=self.PREFIX) as td:
      tempdir = td
      # Show the temp directory exists and is empty.
      self.assertTrue(os.path.isdir(tempdir))
      self.assertEquals(os.listdir(tempdir), [])

    # Show the temp directory no longer exists.
    self.assertNotExists(tempdir)

  def testBasicSuccessNotEmpty(self):
    """Test we cleanup tempdir with stuff in it."""
    with osutils.TempDir(prefix=self.PREFIX) as td:
      tempdir = td
      # Show the temp directory exists and is empty.
      self.assertTrue(os.path.isdir(tempdir))
      self.assertEquals(os.listdir(tempdir), [])

      # Create an empty file.
      osutils.Touch(os.path.join(tempdir, 'foo.txt'))

      # Create nested sub directories.
      subdir = os.path.join(tempdir, 'foo', 'bar', 'taco')
      os.makedirs(subdir)
      osutils.Touch(os.path.join(subdir, 'sauce.txt'))

    # Show the temp directory no longer exists.
    self.assertNotExists(tempdir)

  def testErrorCleanup(self):
    """Test we cleanup, even if an exception is raised."""
    try:
      with osutils.TempDir(prefix=self.PREFIX) as td:
        tempdir = td
        raise TempDirTests.HelperException()
    except TempDirTests.HelperException:
      pass

    # Show the temp directory no longer exists.
    self.assertNotExists(tempdir)

  def testCleanupExceptionContextException(self):
    """Test an exception during cleanup if the context DID raise."""
    was_raised = False

    with mock.patch.object(osutils.TempDir, 'Cleanup',
                           side_effect=TempDirTests.HelperException):
      try:
        with osutils.TempDir(prefix=self.PREFIX) as td:
          tempdir = td
          raise TempDirTests.HelperExceptionInner()
      except TempDirTests.HelperExceptionInner:
        was_raised = True

    # Show that the exception exited the context.
    self.assertTrue(was_raised)

    # Cleanup the dir leaked by our mock exception.
    os.rmdir(tempdir)

  def testCleanupExceptionNoContextException(self):
    """Test an exception during cleanup if the context did NOT raise."""
    was_raised = False

    with mock.patch.object(osutils.TempDir, 'Cleanup',
                           side_effect=TempDirTests.HelperException):
      try:
        with osutils.TempDir(prefix=self.PREFIX) as td:
          tempdir = td
      except TempDirTests.HelperException:
        was_raised = True

    # Show that the exception exited the context.
    self.assertTrue(was_raised)

    # Cleanup the dir leaked by our mock exception.
    os.rmdir(tempdir)


class MountTests(cros_test_lib.TestCase):
  """Unittests for osutils mounting and umounting helpers."""

  def testMountTmpfsDir(self):
    """Verify mounting a tmpfs works"""
    cleaned = False
    with osutils.TempDir(prefix='chromite.test.osutils') as tempdir:
      st_before = os.stat(tempdir)
      try:
        # Mount the dir and verify it worked.
        osutils.MountTmpfsDir(tempdir)
        st_after = os.stat(tempdir)
        self.assertNotEqual(st_before.st_dev, st_after.st_dev)

        # Unmount the dir and verify it worked.
        osutils.UmountDir(tempdir)
        cleaned = True

        # Finally make sure it's cleaned up.
        self.assertFalse(os.path.exists(tempdir))
      finally:
        if not cleaned:
          cros_build_lib.SudoRunCommand(['umount', '-lf', tempdir],
                                        error_code_ok=True)


class IteratePathParentsTest(cros_test_lib.TestCase):
  """Test parent directory iteration functionality."""

  def _RunForPath(self, path, expected):
    result_components = []
    for p in osutils.IteratePathParents(path):
      result_components.append(os.path.basename(p))

    result_components.reverse()
    if expected is not None:
      self.assertEquals(expected, result_components)

  def testIt(self):
    """Run the test vectors."""
    vectors = {
        '/': [''],
        '/path/to/nowhere': ['', 'path', 'to', 'nowhere'],
        '/path/./to': ['', 'path', 'to'],
        '//path/to': ['', 'path', 'to'],
        'path/to': None,
        '': None,
    }
    for p, e in vectors.iteritems():
      self._RunForPath(p, e)


class FindInPathParentsTest(cros_test_lib.TempDirTestCase):
  """Test FindInPathParents functionality."""

  D = cros_test_lib.Directory

  DIR_STRUCT = [
    D('a', [
      D('.repo', []),
      D('b', [
        D('c', [])
      ])
    ])
  ]

  START_PATH = os.path.join('a', 'b', 'c')

  def setUp(self):
    cros_test_lib.CreateOnDiskHierarchy(self.tempdir, self.DIR_STRUCT)

  def testFound(self):
    """Target is found."""
    found = osutils.FindInPathParents(
      '.repo', os.path.join(self.tempdir, self.START_PATH))
    self.assertEquals(found, os.path.join(self.tempdir, 'a', '.repo'))

  def testNotFound(self):
    """Target is not found."""
    found = osutils.FindInPathParents(
      'does.not/exist', os.path.join(self.tempdir, self.START_PATH))
    self.assertEquals(found, None)


class SourceEnvironmentTest(cros_test_lib.TempDirTestCase):
  """Test ostuil's environmental variable related methods."""

  ENV_WHITELIST = {
      'ENV1': 'monkeys like bananas',
      'ENV3': 'merci',
      'ENV6': '',
  }

  ENV_OTHER = {
      'ENV2': 'bananas are yellow',
      'ENV4': 'de rien',
  }

  ENV = """
declare -x ENV1="monkeys like bananas"
declare -x ENV2="bananas are yellow"
declare -x ENV3="merci"
declare -x ENV4="de rien"
declare -x ENV6=''
declare -x ENVA=('a b c' 'd' 'e 1234 %')
"""

  def setUp(self):
    self.env_file = os.path.join(self.tempdir, 'environment')
    osutils.WriteFile(self.env_file, self.ENV)

  def testWhiteList(self):
    env_dict = osutils.SourceEnvironment(
        self.env_file, ('ENV1', 'ENV3', 'ENV5', 'ENV6'))
    self.assertEquals(env_dict, self.ENV_WHITELIST)

  def testArrays(self):
    env_dict = osutils.SourceEnvironment(self.env_file, ('ENVA',))
    self.assertEquals(env_dict, {'ENVA': 'a b c,d,e 1234 %'})

    env_dict = osutils.SourceEnvironment(self.env_file, ('ENVA',), ifs=' ')
    self.assertEquals(env_dict, {'ENVA': 'a b c d e 1234 %'})


class DeviceInfoTests(cros_build_lib_unittest.RunCommandTestCase):
  """Tests methods retrieving information about devices."""

  FULL_OUTPUT = """
NAME="sda" RM="0" TYPE="disk" SIZE="128G"
NAME="sda1" RM="1" TYPE="part" SIZE="100G"
NAME="sda2" RM="1" TYPE="part" SIZE="28G"
NAME="sdc" RM="1" TYPE="disk" SIZE="7.4G"
NAME="sdc1" RM="1" TYPE="part" SIZE="1G"
NAME="sdc2" RM="1" TYPE="part" SIZE="6.4G"
"""

  PARTIAL_OUTPUT = """
NAME="sdc" RM="1" TYPE="disk" SIZE="7.4G"
NAME="sdc1" RM="1" TYPE="part" SIZE="1G"
NAME="sdc2" RM="1" TYPE="part" SIZE="6.4G"
"""

  def testListBlockDevices(self):
    """Tests that we can list al block devices correctly."""
    self.rc.AddCmdResult(partial_mock.Ignore(), output=self.FULL_OUTPUT)
    devices = osutils.ListBlockDevices()
    self.assertEqual(devices[0].NAME, 'sda')
    self.assertEqual(devices[0].RM, '0')
    self.assertEqual(devices[0].TYPE, 'disk')
    self.assertEqual(devices[0].SIZE, '128G')
    self.assertEqual(devices[3].NAME, 'sdc')
    self.assertEqual(devices[3].RM, '1')
    self.assertEqual(devices[3].TYPE, 'disk')
    self.assertEqual(devices[3].SIZE, '7.4G')

  def testGetDeviceSize(self):
    """Tests that we can get the size of a device."""
    self.rc.AddCmdResult(partial_mock.Ignore(), output=self.PARTIAL_OUTPUT)
    self.assertEqual(osutils.GetDeviceSize('/dev/sdc'), '7.4G')


if __name__ == '__main__':
  cros_test_lib.main()