Add command line

Author: xadupre

_unittests/ut_xrun_doc/data/two_nodes.onnx

@@ -0,0 +1,39 @@
+import os
+import unittest
+from contextlib import redirect_stdout
+from io import StringIO
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic._command_lines_parser import main
+
+
+class TestCommandLines(ExtTestCase):
+    @property
+    def dummy_path(self):
+        return os.path.abspath(
+            os.path.join(os.path.dirname(__file__), "data", "two_nodes.onnx")
+        )
+
+    def test_parser_print(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            main(["print", "raw", self.dummy_path])
+        text = st.getvalue()
+        self.assertIn("Add", text)
+
+    def test_parser_find(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            main(["find", "-i", self.dummy_path, "-n", "node_Add_188"])
+        text = st.getvalue()
+        self.assertIsInstance(text, str)
+
+    def test_parser_config(self):
+        st = StringIO()
+        with redirect_stdout(st):
+            main(["config", "-m", "arnir0/Tiny-LLM"])
+        text = st.getvalue()
+        self.assertIn("LlamaForCausalLM", text)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_xrun_doc/test_command_lines_exe.py

@@ -7,33 +7,6 @@
 from textwrap import dedent
 
 
-def get_main_parser() -> ArgumentParser:
-    parser = ArgumentParser(
-        prog="onnx_diagnostic",
-        description="onnx_diagnostic main command line.\n",
-        formatter_class=RawTextHelpFormatter,
-        epilog=textwrap.dedent(
-            """
-        Type 'python -m onnx_diagnostic <cmd> --help'
-        to get help for a specific command.
-
-        lighten    - makes an onnx model lighter by removing the weights,
-        unlighten  - restores an onnx model produces by the previous experiment
-        run        - runs a model and measure the inference time
-        print      - prints the model on standard output
-        find       - find node consuming or producing a result
-        config     - prints a configuration for a model id
-        """
-        ),
-    )
-    parser.add_argument(
-        "cmd",
-        choices=["lighten", "unlighten", "run", "print", "find", "config"],
-        help="Selects a command.",
-    )
-    return parser
-
-
 def get_parser_lighten() -> ArgumentParser:
     parser = ArgumentParser(
         prog="lighten",
@@ -254,6 +227,32 @@ def _cmd_config(argv: List[Any]):
         print(f"task: {task_from_id(args.mid)}")
 
 
+def get_main_parser() -> ArgumentParser:
+    parser = ArgumentParser(
+        prog="onnx_diagnostic",
+        description="onnx_diagnostic main command line.\n",
+        formatter_class=RawTextHelpFormatter,
+        epilog=textwrap.dedent(
+            """
+        Type 'python -m onnx_diagnostic <cmd> --help'
+        to get help for a specific command.
+
+        lighten    - makes an onnx model lighter by removing the weights,
+        unlighten  - restores an onnx model produces by the previous experiment
+        print      - prints the model on standard output
+        find       - find node consuming or producing a result
+        config     - prints a configuration for a model id
+        """
+        ),
+    )
+    parser.add_argument(
+        "cmd",
+        choices=["lighten", "unlighten", "run", "print", "find", "config"],
+        help="Selects a command.",
+    )
+    return parser
+
+
 def main(argv: Optional[List[Any]] = None):
     fcts = dict(
         lighten=_cmd_lighten,