Files
dvc/tests/unit/command/test_debug.py
T

39 lines
1.1 KiB
Python

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"]())