test_dot.py 文件源码

def test_filenames_and_formats(model):
    ipydisp = pytest.importorskip('IPython.display')

    # Test with a variety of user provided args
    filenames = ['modelpdf', 'model.pdf', 'model.pdf', 'modelpdf',
                 'model.pdf.svg']
    formats = ['svg', None, 'svg', None, None]
    targets = ['modelpdf.svg', 'model.pdf', 'model.pdf.svg', 'modelpdf.png',
               'model.pdf.svg']

    result_types = {
        'png': ipydisp.Image,
        'jpeg': ipydisp.Image,
        'dot': type(None),
        'pdf': type(None),
        'svg': ipydisp.SVG,
    }

    for filename, format, target in zip(filenames, formats, targets):
        expected_result_type = result_types[target.split('.')[-1]]
        result = dot_graph(model, filename=filename, format=format)
        assert os.path.isfile(target)
        assert isinstance(result, expected_result_type)
        _ensure_not_exists(target)