tests: Add test_debug

dvc/_debug.py

@@ -12,13 +12,13 @@ def viztracer_profile(
     max_stack_depth: int = -1,
 ):
     try:
-        from viztracer import VizTracer  # pylint: disable=import-error
+        import viztracer  # pylint: disable=import-error
     except ImportError:
         print("Failed to run profiler, viztracer is not installed")
         yield
         return
 
-    tracer = VizTracer(max_stack_depth=max_stack_depth)
+    tracer = viztracer.VizTracer(max_stack_depth=max_stack_depth)
 
     tracer.start()
     yield
@@ -126,12 +126,12 @@ def debugtools(args: "Namespace" = None, **kwargs):
             stack.enter_context(
                 yappi_profile(path=lambda: output.format(datetime.now()))
             )
-        if kw.get("viztracer"):
+        if kw.get("viztracer") or kw.get("viztracer_depth"):
             output = "viztracer.dvc-{0:%Y%m%d}_{0:%H%M%S}.json"
             stack.enter_context(
                 viztracer_profile(
                     path=lambda: output.format(datetime.now()),
-                    max_stack_depth=kw.get("max_stack_depth", -1),
+                    max_stack_depth=kw.get("viztracer_depth") or -1,
                 )
             )
         yield
@@ -149,7 +149,9 @@ def add_debugging_flags(parser):
     parser.add_argument(
         "--viztracer", action="store_true", default=False, help=SUPPRESS
     )
-    parser.add_argument("--max-stack-depth", default=-1, help=SUPPRESS)
+    parser.add_argument(
+        "--viztracer-depth", type=int, default=-1, help=SUPPRESS
+    )
     parser.add_argument("--cprofile-dump", help=SUPPRESS)
     parser.add_argument(
         "--pdb", action="store_true", default=False, help=SUPPRESS

tests/unit/command/test_debug.py

@@ -0,0 +1,38 @@
+from dvc.cli import main
+
+
+def test_viztracer(tmp_dir, dvc, mocker):
+    viztracer_profile = mocker.patch("dvc._debug.viztracer_profile")
+
+    assert main(["status", "--viztracer"]) == 0
+
+    args = viztracer_profile.call_args[1]
+    assert callable(args["path"])
+    assert args["max_stack_depth"] == -1
+
+    assert main(["status", "--viztracer", "--viztracer-depth", "5"]) == 0
+    args = viztracer_profile.call_args[1]
+    assert args["max_stack_depth"] == 5
+
+    assert main(["status", "--viztracer-depth", "2"]) == 0
+    args = viztracer_profile.call_args[1]
+    assert args["max_stack_depth"] == 2
+
+
+def test_viztracer_internal(tmp_dir, dvc, mocker):
+    import sys
+
+    viztracer = mocker.MagicMock()
+    sys.modules["viztracer"] = viztracer
+    from dvc import _debug
+
+    profile = mocker.spy(_debug, "viztracer_profile")
+    tracer = viztracer.VizTracer.return_value
+    assert main(["status", "--viztracer"]) == 0
+
+    args = viztracer.VizTracer.call_args[1]
+    assert args["max_stack_depth"] == -1
+
+    tracer.start.assert_called_once()
+    tracer.stop.assert_called_once()
+    tracer.save.assert_called_once_with(profile.call_args[1]["path"]())