improve test case and wrap code in function

ourownstory · Nov 22, 2022 · c7e42d2 · c7e42d2
1 parent 1d2bc34
commit c7e42d2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/neuralprophet/__main__.py b/neuralprophet/__main__.py
@@ -5,9 +5,12 @@
 
 from neuralprophet._version import __version__
 
-parser = argparse.ArgumentParser(description="NeuralProphet")
-parser.add_argument("-v", "--version", action="version", version="%(prog)s " + __version__)
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser(description="NeuralProphet")
+    parser.add_argument("-v", "--version", action="version", version="%(prog)s " + __version__)
+    return parser.parse_args(args)
 
 
 if __name__ == "__main__":
-    parser.parse_args()
+    parse_args()
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -1,13 +1,13 @@
 import pytest
 
-from neuralprophet.__main__ import parser
+from neuralprophet.__main__ import parse_args
 from neuralprophet._version import __version__
 
 
 def test_main_file(capsys):
     with pytest.raises(SystemExit) as exit_info:
-        parser.parse_args(["--version"])
+        parse_args(["--version"])
 
     out, _ = capsys.readouterr()
-    assert "<ExceptionInfo SystemExit(0) tblen=8>" == str(exit_info)
+    assert exit_info.value.code == 0
     assert __version__ in out