wrapper3: add unittests

This module has some complicated logic which begs for unittests.

BUG=chromium:1170007
TEST=CQ passes

Change-Id: I080aa2599a049f533f70258c7ba4f21099489afa
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2692518
Tested-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Chris McDonald <cjmcdonald@chromium.org>
diff --git a/scripts/wrapper3_unittest.py b/scripts/wrapper3_unittest.py
new file mode 100644
index 0000000..f6aff6c
--- /dev/null
+++ b/scripts/wrapper3_unittest.py
@@ -0,0 +1,189 @@
+# 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.
+
+"""Unittests for wrapper3"""
+
+import os
+from pathlib import Path
+import sys
+from typing import List, Union
+
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import timeout_util
+from chromite.scripts import wrapper3
+
+
+assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'
+
+
+WRAPPER = Path(__file__).resolve().parent / 'wrapper3.py'
+
+
+class FindTargetTests(cros_test_lib.TempDirTestCase):
+  """Tests for FindTarget()."""
+
+  def setUp(self):
+    # It is important that the code we test is not peturbed because the funky
+    # path logic relies heavily upon it to make decisions.
+    self.assertIsNone(wrapper3.CHROMITE_PATH)
+
+    # TODO)vapier): Switch tempdir to pathlib.
+    self.tempdir = Path(self.tempdir)
+
+    # Create a skeleton chromite layout.
+    # tmpdir/
+    #   chromite/
+    #     bin/<wrapper>
+    #     scripts/
+    #     api -> <real chromite>/api/
+    #     lib -> <real chromite>/lib/
+    #     utils -> <real chromite>/utils/
+    #     __init__.py     # Marker file for Python module import.
+    #     PRESUBMIT.cfg   # Marker file for our wrapper to find chromite.
+    self.chromite_dir = self.tempdir / 'chromite'
+    self.bindir = self.chromite_dir / 'bin'
+    self.bindir.mkdir(parents=True)
+    self.scripts_dir = self.chromite_dir / 'scripts'
+    self.scripts_dir.mkdir()
+    for subdir in ('api', 'lib', 'utils'):
+      (self.chromite_dir / subdir).symlink_to(
+          Path(constants.CHROMITE_DIR) / subdir)
+    for subfile in ('__init__.py', 'PRESUBMIT.cfg'):
+      (self.chromite_dir / subfile).touch()
+    self.wrapper = self.scripts_dir / WRAPPER.name
+    # Copy over the wrapper.  We can't just symlink it because the code also
+    # walks & resolves symlinks on itself.  Try hardlink at first, but if the
+    # tempdir is on a diff mount, fallback to a copy.
+    try:
+      if sys.version_info >= (3, 8):
+        self.wrapper.link_to(WRAPPER)
+      else:
+        os.link(WRAPPER, self.wrapper)
+    except OSError:
+      self.wrapper.write_bytes(WRAPPER.read_bytes())
+      self.wrapper.chmod(0o755)
+
+  @staticmethod
+  def insert_path(var: str, value: str):
+    """Insert |value| into the start of the environment |var|."""
+    if var in os.environ:
+      value += f':{os.environ[var]}'
+    os.environ[var] = value
+
+  def gen_script(self, path: Path, wrapper: Path = None):
+    """Create a script at |path|."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path = path.with_suffix('.py')
+    path.write_text('def main(argv):\n  print("hi", argv)\n')
+    if wrapper is None:
+      wrapper = path.with_suffix('')
+    wrapper.symlink_to(self.wrapper)
+
+  def run_script(self, argv: List[Union[Path, str]], **kwargs):
+    """Run |prog| and return the output."""
+    # Log the directory layout to help with debugging.
+    try:
+      cros_build_lib.run(['tree', '-p', str(self.tempdir)], encoding='utf-8',
+                         print_cmd=False)
+    except cros_build_lib.RunCommandError:
+      pass
+
+    # Helper to include a small timeout in case of bugs.
+    with timeout_util.Timeout(30):
+      return cros_build_lib.run([str(x) for x in argv], capture_output=True,
+                                encoding='utf-8', **kwargs)
+
+  def _run_tests(self, prog: Path, verify=None, **kwargs):
+    """Run |prog| in the different fun ways."""
+    if verify is None:
+      verify = lambda result: self.assertEqual('hi []\n', result.output)
+
+    # Execute absolute path.
+    result = self.run_script([prog], **kwargs)
+    verify(result)
+
+    # Execute ./ relative path.
+    result = self.run_script([f'./{prog.name}'], cwd=prog.parent, **kwargs)
+    verify(result)
+
+    # Execute ./path/ relative path.
+    result = self.run_script([f'./{prog.parent.name}/{prog.name}'],
+                             cwd=prog.parent.parent, **kwargs)
+    verify(result)
+
+    # Run via $PATH.
+    self.insert_path('PATH', str(prog.parent))
+    result = self.run_script([prog.name], **kwargs)
+    verify(result)
+
+  def testExternal(self):
+    """Verify use from outside of chromite/ works with main() scripts."""
+    prog = self.tempdir / 'path' / 'prog'
+    self.gen_script(prog)
+    self._run_tests(prog)
+
+  def testChromiteBin(self):
+    """Verify chromite/bin/ works with module in chromite/scripts/."""
+    prog = self.bindir / 'prog'
+    self.gen_script(self.scripts_dir / prog.name, prog)
+    self._run_tests(prog)
+
+  def testChromiteScripts(self):
+    """Verify chromite/scripts/ works with main() scripts."""
+    prog = self.scripts_dir / 'prog'
+    self.gen_script(prog)
+    self._run_tests(prog)
+
+  def testChromiteCustomdir(self):
+    """Verify chromite/customdir/ works with main() scripts."""
+    prog = self.chromite_dir / 'customdir' / 'prog'
+    self.gen_script(prog)
+    self._run_tests(prog)
+
+  def testChromiteTopdir(self):
+    """Verify chromite/ works with main() scripts."""
+    prog = self.chromite_dir / 'prog'
+    self.gen_script(prog)
+    self._run_tests(prog)
+
+  def testUnittests(self):
+    """Allow direct execution of unittests."""
+    prog = self.chromite_dir / 'subdir' / 'prog_unittest'
+    prog.parent.mkdir(parents=True, exist_ok=True)
+    path = prog.with_suffix('.py')
+    path.write_text('import sys; print("hi", sys.argv[1:])\n')
+    prog.symlink_to(self.wrapper)
+    self._run_tests(prog)
+
+  def testTests(self):
+    """Allow direct execution of tests."""
+    prog = self.chromite_dir / 'subdir' / 'prog_unittest'
+    prog.parent.mkdir(parents=True, exist_ok=True)
+    prog.symlink_to(self.wrapper)
+    prog.with_suffix('.py').write_text(
+        'import sys; print("hi", sys.argv[1:])\n')
+    self._run_tests(prog)
+
+  def testWrapper(self):
+    """Fail quickly when running the wrapper directly."""
+    verify = lambda result: self.assertEqual(result.returncode, 100)
+    self._run_tests(self.wrapper, verify=verify, check=False)
+
+  def testMissingScript(self):
+    """Fail quickly if wrapped script is missing."""
+    verify = lambda result: self.assertNotEqual(result.returncode, 0)
+    prog = self.bindir / 'prog'
+    prog.symlink_to(self.wrapper)
+    self._run_tests(prog, verify=verify, check=False)
+
+  def testBrokenScript(self):
+    """Fail quickly if wrapped script is corrupt."""
+    verify = lambda result: self.assertNotEqual(result.returncode, 0)
+    prog = self.scripts_dir / 'prog'
+    prog.symlink_to(self.wrapper)
+    # Script has syntax errors and cannot be imported.
+    prog.with_suffix('.py').write_text('}')
+    self._run_tests(prog, verify=verify, check=False)