From 7741c94d00c2ea309ffd598056c81ab5d566898a Mon Sep 17 00:00:00 2001
From: Jonatan Schlag <jonatan.schlag@ipfire.org>
Date: Sun, 13 May 2018 19:16:08 +0200
Subject: [PATCH] Add main() function

Fixes: #11719

Signed-off-by: Jonatan Schlag <jonatan.schlag@ipfire.org>
---
 nitsi.in | 67 +++++++++++++++++++++++++++++++-------------------------
 1 file changed, 37 insertions(+), 30 deletions(-)

diff --git a/nitsi.in b/nitsi.in
index f8f7945..f45dbe6 100755
--- a/nitsi.in
+++ b/nitsi.in
@@ -3,6 +3,7 @@
 from nitsi.test import test
 from nitsi.logger import init_logging
 import logging
+import argparse
 
 from nitsi.recipe import RecipeExeption
 
@@ -19,9 +20,7 @@ ch.setFormatter(formatter)
 # add the handlers to the logger
 logger.addHandler(ch)
 
-if __name__ == "__main__":
-    import argparse
-
+def main():
     parser = argparse.ArgumentParser()
 
     parser.add_argument("-d", "--directory", dest="dir")
@@ -31,34 +30,42 @@ if __name__ == "__main__":
 
     args = parser.parse_args()
 
+    # We just log the version and exit
     if args.version:
         logger.info("nitsi version: {}".format("@PACKAGE_VERSION@"))
-    else:
-        log_dir = init_logging(args.dir)
-        # We now going to log everything to log_dir/genaral.log
-        fh = logging.FileHandler("{}/general.log".format(log_dir))
-        fh.setLevel(logging.DEBUG)
-        logger.addHandler(fh)
-        logger.debug("We now logging everything to {}/general.log".format(log_dir))
-        try:
-            currenttest = test(args.dir, log_dir)
-            currenttest.read_settings()
-            currenttest.virtual_environ_setup()
-            currenttest.load_recipe()
-        except RecipeExeption as e:
-            logger.exception(e)
-            exit(2)
+        return 0
+
+    # For all other stuff we need logging to a file
+    log_dir = init_logging(args.dir)
+    # We now going to log everything to log_dir/genaral.log
+    fh = logging.FileHandler("{}/general.log".format(log_dir))
+    fh.setLevel(logging.DEBUG)
+    logger.addHandler(fh)
+    logger.debug("We now logging everything to {}/general.log".format(log_dir))
 
-        try:
-            currenttest.virtual_environ_start()
-            currenttest.run_recipe()
-        except TestException as e:
-            logger.exception(e)
-            exit(1)
-        except BaseException as e:
-            logger.exception(e)
-            exit(3)
-        finally:
-            currenttest.virtual_environ_stop()
+    # here we run a test
+    try:
+        currenttest = test(args.dir, log_dir)
+        currenttest.read_settings()
+        currenttest.virtual_environ_setup()
+        currenttest.load_recipe()
+    except RecipeExeption as e:
+        logger.exception(e)
+        exit(2)
 
-        exit(0)
\ No newline at end of file
+    try:
+        currenttest.virtual_environ_start()
+        currenttest.run_recipe()
+    except TestException as e:
+        logger.exception(e)
+        return 1
+    except BaseException as e:
+        logger.exception(e)
+        return 3
+    finally:
+        currenttest.virtual_environ_stop()
+
+    return 0
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
-- 
2.39.2