test_image: allow people to run specific tests

The current code runs all tests all the time.  This can be painful for
people adding new tests and doing iterative development.  Add support
for running specific unittests like standard python unittest runners.

Also add a --list option so people can quickly discover tests.

BUG=chromium:775932
TEST=test_image & unittests still pass
TEST=running specific tests works

Change-Id: I5822190a91a168de6bd93489a41ac8deaf0d73de
Reviewed-on: https://chromium-review.googlesource.com/726895
Commit-Ready: Mike Frysinger <vapier@chromium.org>
Tested-by: Mike Frysinger <vapier@chromium.org>
Reviewed-by: Jorge Lucangeli Obes <jorgelo@chromium.org>
diff --git a/scripts/test_image.py b/scripts/test_image.py
index a86abe8..5391577 100644
--- a/scripts/test_image.py
+++ b/scripts/test_image.py
@@ -27,6 +27,12 @@
   parser.add_argument('image_dir', type='path',
                       help='Image directory (or file) with mount_image.sh and '
                            'umount_image.sh')
+
+  parser.add_argument('-l', '--list', default=False, action='store_true',
+                      help='List all the available tests')
+  parser.add_argument('tests', nargs='*', metavar='test',
+                      help='Specific tests to run (default runs all)')
+
   opts = parser.parse_args(args)
   opts.Freeze()
   return opts
@@ -66,7 +72,29 @@
   # We use a different prefix here so that unittest DO NOT pick up the
   # image tests automatically because they depend on a proper environment.
   loader.testMethodPrefix = 'Test'
-  all_tests = loader.loadTestsFromName('chromite.cros.test.image_test')
+  tests_namespace = 'chromite.cros.test.image_test'
+  if opts.tests:
+    tests = ['%s.%s' % (tests_namespace, x) for x in opts.tests]
+  else:
+    tests = (tests_namespace,)
+  all_tests = loader.loadTestsFromNames(tests)
+
+  # If they just want to see the lists of tests, show them now.
+  if opts.list:
+    def _WalkSuite(suite):
+      for test in suite:
+        if isinstance(test, unittest.BaseTestSuite):
+          for result in _WalkSuite(test):
+            yield result
+        else:
+          yield (test.id()[len(tests_namespace) + 1:],
+                 test.shortDescription() or '')
+
+    test_list = list(_WalkSuite(all_tests))
+    maxlen = max(len(x[0]) for x in test_list)
+    for name, desc in test_list:
+      print('%-*s  %s' % (maxlen, name, desc))
+    return
 
   # Run them in the image directory.
   runner = image_test_lib.ImageTestRunner()