From 8a67befa1c0231f1fc2f3e019b0dbf6127271910 Mon Sep 17 00:00:00 2001
From: jonathanmetzman <31354670+jonathanmetzman@users.noreply.github.com>
Date: Mon, 22 Mar 2021 10:39:27 -0700
Subject: [infra/helper.py] Unittest shell base-runner-debug (#5470)

Make sure shell-base-runner-debug works.
---
 infra/helper.py | 80 +++++++++++++++++++++++++++++++++------------------------
 1 file changed, 46 insertions(+), 34 deletions(-)

(limited to 'infra/helper.py')

diff --git a/infra/helper.py b/infra/helper.py
index 6048d977..6c2950ab 100755
--- a/infra/helper.py
+++ b/infra/helper.py
@@ -64,12 +64,56 @@ LANGUAGES_WITH_COVERAGE_SUPPORT = ['c', 'c++', 'go', 'rust']
 # pylint: disable=too-many-lines
 
 
-def main():  # pylint: disable=too-many-branches,too-many-return-statements,too-many-statements
+def main():  # pylint: disable=too-many-branches,too-many-return-statements
   """Get subcommand from program arguments and do it."""
   os.chdir(OSS_FUZZ_DIR)
   if not os.path.exists(BUILD_DIR):
     os.mkdir(BUILD_DIR)
 
+  args = parse_args()
+
+  # We have different default values for `sanitizer` depending on the `engine`.
+  # Some commands do not have `sanitizer` argument, so `hasattr` is necessary.
+  if hasattr(args, 'sanitizer') and not args.sanitizer:
+    if args.engine == 'dataflow':
+      args.sanitizer = 'dataflow'
+    else:
+      args.sanitizer = 'address'
+
+  if args.command == 'generate':
+    return generate(args)
+  if args.command == 'build_image':
+    return build_image(args)
+  if args.command == 'build_fuzzers':
+    return build_fuzzers(args)
+  if args.command == 'check_build':
+    return check_build(args)
+  if args.command == 'download_corpora':
+    return download_corpora(args)
+  if args.command == 'run_fuzzer':
+    return run_fuzzer(args)
+  if args.command == 'coverage':
+    return coverage(args)
+  if args.command == 'reproduce':
+    return reproduce(args)
+  if args.command == 'shell':
+    return shell(args)
+  if args.command == 'pull_images':
+    return pull_images(args)
+
+  return 0
+
+
+def parse_args(args=None):
+  """Parses args using argparser and returns parsed args."""
+  # Use default argument None for args so that in production, argparse does its
+  # normal behavior, but unittesting is easier.
+  parser = get_parser()
+  return parser.parse_args(args)
+
+
+def get_parser():  # pylint: disable=too-many-statements
+  """Returns an argparse parser."""
   parser = argparse.ArgumentParser('helper.py', description='oss-fuzz helpers')
   subparsers = parser.add_subparsers(dest='command')
 
@@ -192,39 +236,7 @@ def main():  # pylint: disable=too-many-branches,too-many-return-statements,too-
   _add_environment_args(shell_parser)
 
   subparsers.add_parser('pull_images', help='Pull base images.')
-
-  args = parser.parse_args()
-
-  # We have different default values for `sanitizer` depending on the `engine`.
-  # Some commands do not have `sanitizer` argument, so `hasattr` is necessary.
-  if hasattr(args, 'sanitizer') and not args.sanitizer:
-    if args.engine == 'dataflow':
-      args.sanitizer = 'dataflow'
-    else:
-      args.sanitizer = 'address'
-
-  if args.command == 'generate':
-    return generate(args)
-  if args.command == 'build_image':
-    return build_image(args)
-  if args.command == 'build_fuzzers':
-    return build_fuzzers(args)
-  if args.command == 'check_build':
-    return check_build(args)
-  if args.command == 'download_corpora':
-    return download_corpora(args)
-  if args.command == 'run_fuzzer':
-    return run_fuzzer(args)
-  if args.command == 'coverage':
-    return coverage(args)
-  if args.command == 'reproduce':
-    return reproduce(args)
-  if args.command == 'shell':
-    return shell(args)
-  if args.command == 'pull_images':
-    return pull_images(args)
-
-  return 0
+  return parser
 
 
 def is_base_image(image_name):
-- 
cgit v1.2.3