def test_draw_tree_exists_ax():
# Create decision tree classifier object
clf = DecisionTreeClassifier(random_state=0)
# Train model
clf.fit(iris_x, iris_y)
pyplot.figure()
ax = pyplot.gca()
ax.set_title("My ax")
draw_tree(clf, [
'sepal length (cm)', 'sepal width (cm)', 'petal length (cm)',
'petal width (cm)'
], ['setosa', 'versicolor', 'virginica'],
ax=ax)
result_path = Path(__file__).parents[0].absolute().joinpath(
"result_images").joinpath("test_visualization_aids").joinpath(
"test_draw_tree_exists_ax.png")
pyplot.savefig(str(result_path))
baseline_path = Path(__file__).parents[0].absolute().joinpath(
"baseline_images").joinpath("test_visualization_aids").joinpath(
"test_draw_tree_exists_ax.png")
pyplot.cla()
pyplot.close(pyplot.gcf())
compare_images_from_paths(str(baseline_path), str(result_path))
def test_draw_tree_deprecated():
with pytest.warns(DeprecationWarning):
# Create decision tree classifier object
clf = DecisionTreeClassifier(random_state=0)
# Train model
clf.fit(iris_x, iris_y)
draw_tree(clf, [
'sepal length (cm)', 'sepal width (cm)', 'petal length (cm)',
'petal width (cm)'
], ['setosa', 'versicolor', 'virginica'])