Skip to content

Commit 3ad94d9

Browse files
authored
Figure.savefig: Add a new test for the show parameter and simplify existing tests. (#3568)
* Get rid of try-except in test_figure_savefig_geotiff * Simplify test_figure_savefig_unknown_extension * Simplify test_figure_savefig with unittest.mock * Add one more test for Figure.savefig(show=True)
1 parent f41974d commit 3ad94d9

File tree

1 file changed

+30
-66
lines changed

1 file changed

+30
-66
lines changed

pygmt/tests/test_figure.py

Lines changed: 30 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from pygmt.helpers import GMTTempFile
1818

1919
_HAS_IPYTHON = bool(importlib.util.find_spec("IPython"))
20+
_HAS_RIOXARRAY = bool(importlib.util.find_spec("rioxarray"))
2021

2122

2223
def test_figure_region():
@@ -100,7 +101,7 @@ def test_figure_savefig_geotiff():
100101
geofname = Path("test_figure_savefig_geotiff.tiff")
101102
fig.savefig(geofname)
102103
assert geofname.exists()
103-
# The .pgw should not exist
104+
# The .pgw file should not exist
104105
assert not geofname.with_suffix(".pgw").exists()
105106

106107
# Save as TIFF
@@ -109,7 +110,7 @@ def test_figure_savefig_geotiff():
109110
assert fname.exists()
110111

111112
# Check if a TIFF is georeferenced or not
112-
try:
113+
if _HAS_RIOXARRAY:
113114
import rioxarray
114115
from rasterio.errors import NotGeoreferencedWarning
115116
from rasterio.transform import Affine
@@ -147,8 +148,6 @@ def test_figure_savefig_geotiff():
147148
a=1.0, b=0.0, c=0.0, d=0.0, e=1.0, f=0.0
148149
)
149150
assert len(record) == 1
150-
except ImportError:
151-
pass
152151
geofname.unlink()
153152
fname.unlink()
154153

@@ -170,9 +169,7 @@ def test_figure_savefig_unknown_extension():
170169
"""
171170
fig = Figure()
172171
fig.basemap(region="10/70/-300/800", projection="X3i/5i", frame="af")
173-
prefix = "test_figure_savefig_unknown_extension"
174-
fmt = "test"
175-
fname = f"{prefix}.{fmt}"
172+
fname = "test_figure_savefig_unknown_extension.test"
176173
with pytest.raises(GMTInvalidInput, match="Unknown extension '.test'."):
177174
fig.savefig(fname)
178175

@@ -223,69 +220,23 @@ def test_figure_savefig():
223220
"""
224221
Check if the arguments being passed to psconvert are correct.
225222
"""
226-
kwargs_saved = []
227-
228-
def mock_psconvert(*args, **kwargs): # noqa: ARG001
229-
"""
230-
Just record the arguments.
231-
"""
232-
kwargs_saved.append(kwargs)
233-
234-
fig = Figure()
235-
fig.psconvert = mock_psconvert
236-
237223
prefix = "test_figure_savefig"
238-
239-
fname = f"{prefix}.png"
240-
fig.savefig(fname)
241-
assert kwargs_saved[-1] == {
242-
"prefix": prefix,
243-
"fmt": "g",
244-
"crop": True,
245-
"Qt": 2,
246-
"Qg": 2,
247-
}
248-
249-
fname = f"{prefix}.pdf"
250-
fig.savefig(fname)
251-
assert kwargs_saved[-1] == {
252-
"prefix": prefix,
253-
"fmt": "f",
254-
"crop": True,
255-
"Qt": 2,
256-
"Qg": 2,
224+
common_kwargs = {"prefix": prefix, "crop": True, "Qt": 2, "Qg": 2}
225+
expected_kwargs = {
226+
"png": {"fmt": "g", **common_kwargs},
227+
"pdf": {"fmt": "f", **common_kwargs},
228+
"eps": {"fmt": "e", **common_kwargs},
229+
"kml": {"fmt": "g", "W": "+k", **common_kwargs},
257230
}
258231

259-
fname = f"{prefix}.png"
260-
fig.savefig(fname, transparent=True)
261-
assert kwargs_saved[-1] == {
262-
"prefix": prefix,
263-
"fmt": "G",
264-
"crop": True,
265-
"Qt": 2,
266-
"Qg": 2,
267-
}
268-
269-
fname = f"{prefix}.eps"
270-
fig.savefig(fname)
271-
assert kwargs_saved[-1] == {
272-
"prefix": prefix,
273-
"fmt": "e",
274-
"crop": True,
275-
"Qt": 2,
276-
"Qg": 2,
277-
}
232+
with patch.object(Figure, "psconvert") as mock_psconvert:
233+
fig = Figure()
234+
for fmt, expected in expected_kwargs.items():
235+
fig.savefig(f"{prefix}.{fmt}")
236+
mock_psconvert.assert_called_with(**expected)
278237

279-
fname = f"{prefix}.kml"
280-
fig.savefig(fname)
281-
assert kwargs_saved[-1] == {
282-
"prefix": prefix,
283-
"fmt": "g",
284-
"crop": True,
285-
"Qt": 2,
286-
"Qg": 2,
287-
"W": "+k",
288-
}
238+
fig.savefig(f"{prefix}.png", transparent=True)
239+
mock_psconvert.assert_called_with(fmt="G", **common_kwargs)
289240

290241

291242
def test_figure_savefig_worldfile():
@@ -309,6 +260,19 @@ def test_figure_savefig_worldfile():
309260
fig.savefig(fname=imgfile.name, worldfile=True)
310261

311262

263+
def test_figure_savefig_show():
264+
"""
265+
Check if the external viewer is launched when the show parameter is specified.
266+
"""
267+
fig = Figure()
268+
fig.basemap(region=[0, 1, 0, 1], projection="X1c/1c", frame=True)
269+
prefix = "test_figure_savefig_show"
270+
with patch("pygmt.figure.launch_external_viewer") as mock_viewer:
271+
with GMTTempFile(prefix=prefix, suffix=".png") as imgfile:
272+
fig.savefig(imgfile.name, show=True)
273+
assert mock_viewer.call_count == 1
274+
275+
312276
@pytest.mark.skipif(not _HAS_IPYTHON, reason="run when IPython is installed")
313277
def test_figure_show():
314278
"""

0 commit comments

Comments
 (0)