From e424405f47c1a76bb1b0434646b2e8c89da10079 Mon Sep 17 00:00:00 2001 From: Lai Wei Date: Mon, 16 Apr 2018 11:42:31 -0700 Subject: [PATCH] skip unittest if graphviz could not be imported --- tests/python/unittest/test_viz.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/python/unittest/test_viz.py b/tests/python/unittest/test_viz.py index 32a63976afb1..48fea71fb38c 100644 --- a/tests/python/unittest/test_viz.py +++ b/tests/python/unittest/test_viz.py @@ -17,6 +17,7 @@ import mxnet as mx import warnings +import unittest def test_print_summary(): data = mx.sym.Variable('data') @@ -33,6 +34,15 @@ def test_print_summary(): shape["data"]=(1,3,28,28) mx.viz.print_summary(sc1, shape) +def graphviz_exists(): + try: + import graphviz + except ImportError: + return False + else: + return True + +@unittest.skipIf(not graphviz_exists(), "Skip test_plot_network as Graphviz could not be imported") def test_plot_network(): # Test warnings for cyclic graph net = mx.sym.Variable('data')