From 0d5e4ae00622a6f73ee4596afed8a51373839b04 Mon Sep 17 00:00:00 2001
From: Chris Sosa <sosa@chromium.org>
Date: Wed, 30 Mar 2011 11:22:17 -0700
Subject: [PATCH] Add better support for getting the path to crosutils.

Change-Id: I83a56ee38d8d666d3b71b4dc323e26076fab7d96

BUG=chromium-os:13498
TEST=Ran unittests.

Review URL: http://codereview.chromium.org/6765001
---
 lib/cros_build_lib.py          | 39 ++++++++++++++++++++++--
 lib/cros_build_lib_unittest.py | 54 +++++++++++++++++++++++++++++++++-
 2 files changed, 90 insertions(+), 3 deletions(-)

diff --git a/lib/cros_build_lib.py b/lib/cros_build_lib.py
index d846757677..6f331a0b3b 100644
--- a/lib/cros_build_lib.py
+++ b/lib/cros_build_lib.py
@@ -11,8 +11,6 @@ import subprocess
 import sys
 
 _STDOUT_IS_TTY = hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
-CROSUTILS_DIRECTORY = os.path.dirname(os.path.dirname(
-    os.path.realpath(__file__)))
 
 # TODO(sosa):  Move logging to logging module.
 
@@ -326,3 +324,40 @@ def UnmountImage(root_dir, stateful_dir):
               '--stateful_mountpt=%s' % stateful_dir,
              ], print_cmd=False, redirect_stdout=True, redirect_stderr=True,
              cwd=CROSUTILS_DIRECTORY)
+
+
+def GetCrosUtilsPath(source_dir_path=True):
+  """Return the path to src/scripts.
+
+  Args:
+    source_dir_path:  If True, returns the path from the source code directory.
+  """
+  if IsInsideChroot():
+    if source_dir_path:
+      return os.path.join(os.getenv('HOME'), 'trunk', 'src', 'scripts')
+
+    return os.path.join('/usr/lib/crosutils')
+
+  # Outside the chroot => from_source.
+  return os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
+
+
+def GetCrosUtilsBinPath(source_dir_path=True):
+  """Return the path to crosutils/bin.
+
+  Args:
+    source_dir_path:  If True, returns the path from the source code directory.
+  """
+  if IsInsideChroot() and not source_dir_path:
+      return '/usr/bin'
+
+  return os.path.join(GetCrosUtilsPath(source_dir_path), 'bin')
+
+
+def IsInsideChroot():
+  """Returns True if we are inside chroot."""
+  return os.path.exists('/etc/debian_chroot')
+
+
+# TODO(sosa): Remove once all callers use method.
+CROSUTILS_DIRECTORY = GetCrosUtilsPath(True)
diff --git a/lib/cros_build_lib_unittest.py b/lib/cros_build_lib_unittest.py
index 293bff15ab..61268fcfc2 100755
--- a/lib/cros_build_lib_unittest.py
+++ b/lib/cros_build_lib_unittest.py
@@ -6,13 +6,14 @@
 
 """Unit tests for cros_build_lib."""
 
+import mox
 import os
 import tempfile
 import unittest
 
 import cros_build_lib
 
-class CrosBuildLibTest(unittest.TestCase):
+class CrosBuildLibTest(mox.MoxTestBase):
   """Test class for cros_build_lib."""
 
   def testRunCommandSimple(self):
@@ -103,6 +104,57 @@ class CrosBuildLibTest(unittest.TestCase):
     log_fh.close()
     os.remove(log_file)
 
+  def testGetCrosUtilsPathInChroot(self):
+    """Tests whether we can get crosutils from chroot."""
+    self.mox.StubOutWithMock(cros_build_lib, 'IsInsideChroot')
+    crosutils_path_src = '/home/' + os.getenv('USER') + 'trunk/src/scripts'
+    crosutils_path_installed = '/usr/lib/crosutils'
+
+    cros_build_lib.IsInsideChroot().MultipleTimes().AndReturn(True)
+
+    self.mox.ReplayAll()
+    self.assertTrue(cros_build_lib.GetCrosUtilsPath(source_dir_path=True),
+                    crosutils_path_src)
+    self.assertTrue(cros_build_lib.GetCrosUtilsPath(source_dir_path=False),
+                    crosutils_path_installed)
+    self.mox.VerifyAll()
+
+  def testGetCrosUtilsPathOutsideChroot(self):
+    """Tests whether we can get crosutils from outside chroot."""
+    self.mox.StubOutWithMock(cros_build_lib, 'IsInsideChroot')
+    path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..')
+    cros_build_lib.IsInsideChroot().MultipleTimes().AndReturn(False)
+
+    self.mox.ReplayAll()
+    self.assertTrue(cros_build_lib.GetCrosUtilsPath(), path)
+    self.mox.VerifyAll()
+
+  def testGetCrosUtilsBinPath(self):
+    """Tests whether we can get crosutilsbin correctly."""
+    self.mox.StubOutWithMock(cros_build_lib, 'IsInsideChroot')
+    self.mox.StubOutWithMock(cros_build_lib, 'GetCrosUtilsPath')
+    src_path = '/fake/src'
+    chroot_src_path = '/chroot/fake/src'
+    chroot_path = '/usr/bin'
+
+    cros_build_lib.IsInsideChroot().AndReturn(False)
+    cros_build_lib.GetCrosUtilsPath(True).AndReturn(src_path)
+    cros_build_lib.IsInsideChroot().AndReturn(True)
+    cros_build_lib.GetCrosUtilsPath(True).AndReturn(chroot_src_path)
+    cros_build_lib.IsInsideChroot().AndReturn(True)
+
+    self.mox.ReplayAll()
+    # Outside chroot.
+    self.assertTrue(cros_build_lib.GetCrosUtilsBinPath(source_dir_path=True),
+                    src_path + '/bin')
+    # Rest inside chroot.
+    self.assertTrue(cros_build_lib.GetCrosUtilsBinPath(source_dir_path=True),
+                    chroot_src_path + '/bin')
+    self.assertTrue(cros_build_lib.GetCrosUtilsBinPath(source_dir_path=False),
+                    chroot_path)
+    self.mox.VerifyAll()
+
+
 
 if __name__ == '__main__':
   unittest.main()