[iwyu_tool] Refactor for testability and add unit test

- Add execute function to run all invocations (asynchronously)
- Move process invocation to Invocation.run
- Make Invocation.from_compile_command a class method to make it easier
  to stub Invocation.run when under test.

With this in place, add basic test cases for compile command parsing and
execution.

2 files changed, 136 insertions, 28 deletions

diff --git a/iwyu_tool.py b/iwyu_tool.py
index 9445293..1e34835 100755
--- a/iwyu_tool.py
+++ b/iwyu_tool.py
@@ -122,8 +122,8 @@ class Invocation(object):
     def __str__(self):
         return ' '.join(self.command)
 
-    @staticmethod
-    def from_compile_command(entry, extra_args):
+    @classmethod
+    def from_compile_command(cls, entry, extra_args):
         """ Parse a JSON compilation database entry into new Invocation. """
         if 'arguments' in entry:
             # arguments is a command-line in list form.
@@ -141,19 +141,19 @@ class Invocation(object):
             extra_args = ['--driver-mode=cl'] + extra_args
 
         command = [IWYU_EXECUTABLE] + extra_args + compile_args
-        return Invocation(command, entry['directory'])
+        return cls(command, entry['directory'])
 
+    def run(self, verbose):
+        """ Run invocation and collect output. """
+        if verbose:
+            print('# %s' % self)
 
-def run_iwyu(invocation, verbose):
-    """ Run invocation and collect output. """
-    if verbose:
-        print('# %s' % invocation)
-
-    process = subprocess.Popen(invocation.command,
-                               cwd=invocation.cwd,
-                               stdout=subprocess.PIPE,
-                               stderr=subprocess.STDOUT)
-    return process.communicate()[0].decode('utf-8')
+        process = subprocess.Popen(
+            self.command,
+            cwd=self.cwd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT)
+        return process.communicate()[0].decode('utf-8')
 
 
 def parse_compilation_db(compilation_db_path):
@@ -198,30 +198,27 @@ def slice_compilation_db(compilation_db, selection):
     return new_db
 
 
-def main(compilation_db_path, source_files, verbose, formatter, jobs, iwyu_args):
-    """ Entry point. """
-
-    try:
-        compilation_db = parse_compilation_db(compilation_db_path)
-    except IOError as why:
-        print('Failed to parse JSON compilation database: %s' % why)
-        return 1
+def _run_iwyu(invocation, verbose):
+    """ Start the process described by invocation.
 
-    compilation_db = slice_compilation_db(compilation_db, source_files)
+    Only necessary because multiprocessing can't pickle instance methods.
+    """
+    return invocation.run(verbose)
 
-    # Transform compilation db entries into a list of IWYU invocations.
-    invocations = [Invocation.from_compile_command(e, iwyu_args)
-                   for e in compilation_db]
 
+def execute(invocations, verbose, formatter, jobs):
+    """ Execute all invocations with at most `jobs` in parallel. """
     try:
         pool = multiprocessing.Pool(jobs)
         # No actual results in `results`, it's only used for exception handling.
         # Details here: https://stackoverflow.com/a/28660669.
         results = []
+
         for invocation in invocations:
-            results.append(pool.apply_async(run_iwyu,
-                                            (invocation, verbose),
-                                            callback=lambda x: print(formatter(x))))
+            results.append(
+                pool.apply_async(
+                    _run_iwyu, (invocation, verbose),
+                    callback=lambda x: print(formatter(x))))
         pool.close()
         pool.join()
         for r in results:
@@ -233,6 +230,25 @@ def main(compilation_db_path, source_files, verbose, formatter, jobs, iwyu_args)
     return 0
 
 
+def main(compilation_db_path, source_files, verbose, formatter, jobs,
+         iwyu_args):
+    """ Entry point. """
+
+    try:
+        compilation_db = parse_compilation_db(compilation_db_path)
+    except IOError as why:
+        print('Failed to parse JSON compilation database: %s' % why)
+        return 1
+
+    compilation_db = slice_compilation_db(compilation_db, source_files)
+
+    # Transform compilation db entries into a list of IWYU invocations.
+    invocations = [
+        Invocation.from_compile_command(e, iwyu_args) for e in compilation_db
+    ]
+
+    return execute(invocations, verbose, formatter, jobs)
+
 def _bootstrap():
     """ Parse arguments and dispatch to main(). """
 
diff --git a/iwyu_tool_test.py b/iwyu_tool_test.py
new file mode 100755
index 0000000..740a0c5
--- /dev/null
+++ b/iwyu_tool_test.py
@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+
+##===-------------- iwyu_tool_test.py - test for iwyu_tool.py -------------===##
+#
+#                     The LLVM Compiler Infrastructure
+#
+# This file is distributed under the University of Illinois Open Source
+# License. See LICENSE.TXT for details.
+#
+##===----------------------------------------------------------------------===##
+
+import time
+import random
+import unittest
+import iwyu_tool
+
+try:
+    from cStringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+class MockInvocation(iwyu_tool.Invocation):
+    def __init__(self, command=None, cwd=''):
+        iwyu_tool.Invocation.__init__(self, command or [], cwd)
+        self._will_return = ''
+        self._will_block = 0
+
+    def will_block(self, seconds):
+        self._will_block = seconds
+
+    def will_return(self, content):
+        self._will_return = content
+
+    def run(self, verbose):
+        if self._will_block > 0:
+            time.sleep(self._will_block)
+        return self._will_return
+
+
+class IWYUToolTestBase(unittest.TestCase):
+    def setUp(self):
+        self.stdout_stub = StringIO()
+        iwyu_tool.sys.stdout = self.stdout_stub
+
+    def _execute(self, invocations, verbose=False, formatter=None, jobs=1):
+        formatter = formatter or iwyu_tool.DEFAULT_FORMAT
+        formatter = iwyu_tool.FORMATTERS.get(formatter, formatter)
+        return iwyu_tool.execute(invocations, verbose, formatter, jobs)
+
+
+class IWYUToolTests(IWYUToolTestBase):
+    def test_from_compile_command(self):
+        iwyu_args = ['-foo']
+        invocation = iwyu_tool.Invocation.from_compile_command(
+            {
+                'directory': '/home/user/llvm/build',
+                'command': '/usr/bin/clang++ -Iinclude file.cc',
+                'file': 'file.cc'
+            }, iwyu_args)
+        self.assertEqual(
+            invocation.command,
+            [iwyu_tool.IWYU_EXECUTABLE, '-foo', '-Iinclude', 'file.cc'])
+        self.assertEqual(invocation.cwd, '/home/user/llvm/build')
+
+    def test_invocation(self):
+        invocation = MockInvocation()
+        invocation.will_return('BAR')
+        self.assertEqual(self._execute([invocation]), 0)
+        self.assertEqual(self.stdout_stub.getvalue(), 'BAR\n')
+
+    def test_order_asynchronous(self):
+        invocations = [MockInvocation() for _ in range(100)]
+        for n, invocation in enumerate(invocations):
+            invocation.will_return('BAR%d' % n)
+            invocation.will_block(random.random() / 100)
+        self.assertEqual(self._execute(invocations, jobs=100), 0)
+        self.assertSetEqual(
+            set('BAR%d' % n for n in range(100)),
+            set(self.stdout_stub.getvalue().splitlines()))
+
+    def test_order_synchronous(self):
+        invocations = [MockInvocation() for _ in range(100)]
+        for n, invocation in enumerate(invocations):
+            invocation.will_return('BAR%d' % n)
+            invocation.will_block(random.random() / 100)
+        self.assertEqual(self._execute(invocations, jobs=1), 0)
+        self.assertEqual(['BAR%d' % n for n in range(100)],
+                         self.stdout_stub.getvalue().splitlines())
+
+
+if __name__ == '__main__':
+    unittest.main()