import os

from matplotlib import pyplot as plt

import pytest

from unittest import mock

@pytest.mark.backend("gtk3agg", skip_on_importerror=True)

def test_save_figure_return():

from gi.repository import Gtk

fig, ax = plt.subplots()

ax.imshow([[1]])

with mock.patch("gi.repository.Gtk.FileFilter") as fileFilter:

filt = fileFilter.return_value

filt.get_name.return_value = "Portable Network Graphics"

with mock.patch("gi.repository.Gtk.FileChooserDialog") as dialogChooser:

dialog = dialogChooser.return_value

dialog.get_filter.return_value = filt

dialog.get_filename.return_value = "foobar.png"

dialog.run.return_value = Gtk.ResponseType.OK

fname = fig.canvas.manager.toolbar.save_figure()

os.remove("foobar.png")

assert fname == "foobar.png"

with mock.patch("gi.repository.Gtk.MessageDialog"):

dialog.get_filename.return_value = None

dialog.run.return_value = Gtk.ResponseType.OK

fname = fig.canvas.manager.toolbar.save_figure()

assert fname is None