test_dot.py 文件源码

def test_dot_graph(model, tmpdir):
    ipydisp = pytest.importorskip('IPython.display')

    # Use a name that the shell would interpret specially to ensure that we're
    # not vulnerable to shell injection when interacting with `dot`.
    filename = str(tmpdir.join('$(touch should_not_get_created.txt)'))

    # Map from format extension to expected return type.
    result_types = {
        'png': ipydisp.Image,
        'jpeg': ipydisp.Image,
        'dot': type(None),
        'pdf': type(None),
        'svg': ipydisp.SVG,
    }
    for format in result_types:
        target = '.'.join([filename, format])
        _ensure_not_exists(target)
        try:
            result = dot_graph(model, filename=filename, format=format)

            assert not os.path.exists('should_not_get_created.txt')
            assert os.path.isfile(target)
            assert isinstance(result, result_types[format])
        finally:
            _ensure_not_exists(target)

    # format supported by graphviz but not by IPython
    with pytest.raises(ValueError) as excinfo:
        dot_graph(model, filename=filename, format='ps')
    assert "Unknown format" in str(excinfo.value)