Create build_helper() function

4 files changed, 24 insertions, 8 deletions

diff --git a/dnsviz/commands/graph.py b/dnsviz/commands/graph.py
index 2539e35..54fb4fd 100644
--- a/dnsviz/commands/graph.py
+++ b/dnsviz/commands/graph.py
@@ -408,12 +408,16 @@ class GraphArgHelper:
                 if name not in self.names:
                     self.names[name] = None
 
+def build_helper(logger, cmd, subcmd):
+    arghelper = GraphArgHelper(logger)
+    arghelper.build_parser('%s %s' % (cmd, subcmd))
+    return arghelper
+
 def main(argv):
     try:
         test_pygraphviz()
 
-        arghelper = GraphArgHelper(logger)
-        arghelper.build_parser('%s %s' % (sys.argv[0], argv[0]))
+        arghelper = build_helper(logger, sys.argv[0], argv[0])
         arghelper.parse_args(argv[1:])
         logger.setLevel(logging.WARNING)
 
diff --git a/dnsviz/commands/grok.py b/dnsviz/commands/grok.py
index 398dbc1..6cdac3e 100644
--- a/dnsviz/commands/grok.py
+++ b/dnsviz/commands/grok.py
@@ -401,11 +401,15 @@ class GrokArgHelper:
                 if name not in self.names:
                     self.names[name] = None
 
+def build_helper(logger, cmd, subcmd):
+    arghelper = GrokArgHelper(logger)
+    arghelper.build_parser('%s %s' % (cmd, subcmd))
+    return arghelper
+
 def main(argv):
     try:
 
-        arghelper = GrokArgHelper(logger)
-        arghelper.build_parser('%s %s' % (sys.argv[0], argv[0]))
+        arghelper = build_helper(logger, sys.argv[0], argv[0])
         arghelper.parse_args(argv[1:])
         logger.setLevel(logging.WARNING)
 
diff --git a/dnsviz/commands/print.py b/dnsviz/commands/print.py
index 6eb83b2..6efb1cd 100644
--- a/dnsviz/commands/print.py
+++ b/dnsviz/commands/print.py
@@ -542,12 +542,16 @@ class PrintArgHelper:
                 if name not in self.names:
                     self.names[name] = None
 
+def build_helper(logger, cmd, subcmd):
+    arghelper = PrintArgHelper(logger)
+    arghelper.build_parser('%s %s' % (cmd, subcmd))
+    return arghelper
+
 def main(argv):
     try:
         test_pygraphviz()
 
-        arghelper = PrintArgHelper(logger)
-        arghelper.build_parser('%s %s' % (sys.argv[0], argv[0]))
+        arghelper = build_helper(logger, sys.argv[0], argv[0])
         arghelper.parse_args(argv[1:])
         logger.setLevel(logging.WARNING)
 
diff --git a/dnsviz/commands/probe.py b/dnsviz/commands/probe.py
index bb01c0a..28f7365 100644
--- a/dnsviz/commands/probe.py
+++ b/dnsviz/commands/probe.py
@@ -1446,6 +1446,11 @@ class ArgHelper:
         for zone in self._zones_to_serve:
             zone.serve()
 
+def build_helper(resolver, logger, cmd, subcmd):
+    arghelper = ArgHelper(resolver, logger)
+    arghelper.build_parser('%s %s' % (cmd, subcmd))
+    return arghelper
+
 def main(argv):
     global tm
     global th_factories
@@ -1460,8 +1465,7 @@ def main(argv):
             sys.stderr.write('File %s not found or contains no nameserver entries.\n' % RESOLV_CONF)
             sys.exit(1)
 
-        arghelper = ArgHelper(bootstrap_resolver, logger)
-        arghelper.build_parser('%s %s' % (sys.argv[0], argv[0]))
+        arghelper = build_helper(bootstrap_resolver, logger, sys.argv[0], argv[0])
         arghelper.parse_args(argv[1:])
         logger.setLevel(arghelper.get_log_level())