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'])