Fixup GAE standard test runner sample

appengine/standard/localtesting/runner.py

@@ -1,4 +1,5 @@
-#!/usr/bin/python
+#!/usr/bin/env python2
+
 # Copyright 2015 Google Inc
 #
 #    Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,51 +15,50 @@
 #    limitations under the License.
 
 # [START runner]
-import optparse
-import os
-import sys
-import unittest
+"""App Engine local test runner example.
 
+This program handles properly importing the App Engine SDK so that test modules
+can use google.appengine.* APIs and the Google App Engine testbed.
 
-USAGE = """%prog SDK_PATH TEST_PATH
-Run unit tests for App Engine apps.
+Example invocation:
 
-SDK_PATH    Path to Google Cloud or Google App Engine SDK installation, usually
-            ~/google_cloud_sdk
-TEST_PATH   Path to package containing test modules"""
+    $ python runner.py ~/google-cloud-sdk
+"""
 
+import argparse
+import os
+import sys
+import unittest
 
-def append_or_insert_path(path):
-    """Adds GAE path to system path, or appends it to the google path
-    if that already exists.
 
-    Not all Google packages are inside namespaced packages, which means
-    there might be another named `google` on the path, and simply appending
-    the App Engine SDK to the path will not work since the other package
-    will get discovered and used first. This ugly hack emulates namespace
-    packages by first searching if a `google` package already exists by
-    importing it, and if so appending to its path, otherwise it just
-    inserts it into the import path by itself.
-    """
+def fixup_paths(path):
+    """Adds GAE SDK path to system path and appends it to the google path
+    if that already exists."""
+    # Not all Google packages are inside namespace packages, which means
+    # there might be another non-namespace package named `google` already on
+    # the path and simply appending the App Engine SDK to the path will not
+    # work since the other package will get discovered and used first.
+    # This emulates namespace packages by first searching if a `google` package
+    # exists by importing it, and if so appending to its module search path.
     try:
         import google
         google.__path__.append("{0}/google".format(path))
     except ImportError:
-        sys.path.insert(0, path)
+        pass
 
+    sys.path.insert(0, path)
 
-def main(sdk_path, test_path):
-    # If the sdk path points to a google cloud sdk installation
-    # then we should alter it to point to the GAE platform location.
 
+def main(sdk_path, test_path, test_pattern):
+    # If the SDK path points to a Google Cloud SDK installation
+    # then we should alter it to point to the GAE platform location.
     if os.path.exists(os.path.join(sdk_path, 'platform/google_appengine')):
-        append_or_insert_path(
-            os.path.join(sdk_path, 'platform/google_appengine'))
-    else:
-        append_or_insert_path(sdk_path)
+        sdk_path = os.path.join(sdk_path, 'platform/google_appengine')
+
+    # Make sure google.appengine.* modules are importable.
+    fixup_paths(sdk_path)
 
-    # Ensure that the google.appengine.* packages are available
-    # in tests as well as all bundled third-party packages.
+    # Make sure all bundled third-party packages are available.
     import dev_appserver
     dev_appserver.fix_sys_path()
 
@@ -69,22 +69,34 @@ def main(sdk_path, test_path):
         import appengine_config
         (appengine_config)
     except ImportError:
-        print "Note: unable to import appengine_config."
+        print('Note: unable to import appengine_config.')
 
     # Discover and run tests.
-    suite = unittest.loader.TestLoader().discover(test_path, "*_test.py")
-    unittest.TextTestRunner(verbosity=2).run(suite)
+    suite = unittest.loader.TestLoader().discover(test_path, test_pattern)
+    return unittest.TextTestRunner(verbosity=2).run(suite)
 
 
 if __name__ == '__main__':
-    parser = optparse.OptionParser(USAGE)
-    options, args = parser.parse_args()
-    if len(args) != 2:
-        print 'Error: Exactly 2 arguments required.'
-        parser.print_help()
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        'sdk_path',
+        help='The path to the Google App Engine SDK or the Google Cloud SDK.')
+    parser.add_argument(
+        '--test-path',
+        help='The path to look for tests, defaults to the current directory.',
+        default=os.getcwd())
+    parser.add_argument(
+        '--test-pattern',
+        help='The file pattern for test modules, defaults to *_test.py.',
+        default='*_test.py')
+
+    args = parser.parse_args()
+
+    result = main(args.sdk_path, args.test_path, args.test_pattern)
+
+    if not result.wasSuccessful():
         sys.exit(1)
-    SDK_PATH = args[0]
-    TEST_PATH = args[1]
-    main(SDK_PATH, TEST_PATH)
 
 # [END runner]