From 82b0c736801998d5a995e2dd0f0f4898dc95904e Mon Sep 17 00:00:00 2001
From: Dongdong Tian <seisman.info@gmail.com>
Date: Thu, 31 Oct 2024 17:03:24 +0800
Subject: [PATCH] pygmt.set_display: Improve the test by checking if the
 preview image is opened in the expected way (#3548)

---
 pygmt/figure.py            | 25 +++++++++----------
 pygmt/tests/test_figure.py | 49 +++++++++++++++++++++++++++++++++-----
 2 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/pygmt/figure.py b/pygmt/figure.py
index 12aaa67f722..9c0a7a7a763 100644
--- a/pygmt/figure.py
+++ b/pygmt/figure.py
@@ -453,27 +453,28 @@ def set_display(method: Literal["external", "notebook", "none", None] = None):
     >>> import pygmt
     >>> fig = pygmt.Figure()
     >>> fig.basemap(region=[0, 10, 0, 10], projection="X10c/5c", frame=True)
-    >>> fig.show()  # will display a PNG image in the current notebook
+    >>> fig.show()  # Will display a PNG image in the current notebook
     >>>
-    >>> # set the display method to "external"
+    >>> # Set the display method to "external"
     >>> pygmt.set_display(method="external")  # doctest: +SKIP
-    >>> fig.show()  # will display a PDF image using the default PDF viewer
+    >>> fig.show()  # Will display a PDF image using the default PDF viewer
     >>>
-    >>> # set the display method to "none"
+    >>> # Set the display method to "none"
     >>> pygmt.set_display(method="none")
-    >>> fig.show()  # will not show any image
+    >>> fig.show()  # Will not show any image
     >>>
-    >>> # reset to the default display method
+    >>> # Reset to the default display method
     >>> pygmt.set_display(method=None)
-    >>> fig.show()  # again, will show a PNG image in the current notebook
+    >>> fig.show()  # Again, will show a PNG image in the current notebook
     """
     match method:
         case "external" | "notebook" | "none":
-            SHOW_CONFIG["method"] = method  # type: ignore[assignment]
+            SHOW_CONFIG["method"] = method
         case None:
-            SHOW_CONFIG["method"] = _get_default_display_method()  # type: ignore[assignment]
+            SHOW_CONFIG["method"] = _get_default_display_method()
         case _:
-            raise GMTInvalidInput(
-                f"Invalid display method '{method}'. Valid values are 'external',"
-                "'notebook', 'none' or None."
+            msg = (
+                f"Invalid display method '{method}'. "
+                "Valid values are 'external', 'notebook', 'none' or None."
             )
+            raise GMTInvalidInput(msg)
diff --git a/pygmt/tests/test_figure.py b/pygmt/tests/test_figure.py
index 4cf4a027474..3118f675232 100644
--- a/pygmt/tests/test_figure.py
+++ b/pygmt/tests/test_figure.py
@@ -377,13 +377,50 @@ class TestSetDisplay:
 
     def test_set_display(self):
         """
-        Test if pygmt.set_display updates the SHOW_CONFIG variable correctly.
+        Test if pygmt.set_display updates the SHOW_CONFIG variable correctly and
+        Figure.show opens the preview image in the correct way.
         """
-        default_method = SHOW_CONFIG["method"]  # Current default method
-
-        for method in ("notebook", "external", "none"):
-            set_display(method=method)
-            assert SHOW_CONFIG["method"] == method
+        default_method = SHOW_CONFIG["method"]  # Store the current default method.
+
+        fig = Figure()
+        fig.basemap(region=[0, 3, 6, 9], projection="X1c", frame=True)
+
+        # Test the "notebook" display method.
+        set_display(method="notebook")
+        assert SHOW_CONFIG["method"] == "notebook"
+        if _HAS_IPYTHON:
+            with (
+                patch("IPython.display.display") as mock_display,
+                patch("pygmt.figure.launch_external_viewer") as mock_viewer,
+            ):
+                fig.show()
+                assert mock_viewer.call_count == 0
+                assert mock_display.call_count == 1
+        else:
+            with pytest.raises(GMTError):
+                fig.show()
+
+        # Test the "external" display method
+        set_display(method="external")
+        assert SHOW_CONFIG["method"] == "external"
+        with patch("pygmt.figure.launch_external_viewer") as mock_viewer:
+            fig.show()
+            assert mock_viewer.call_count == 1
+        if _HAS_IPYTHON:
+            with patch("IPython.display.display") as mock_display:
+                fig.show()
+                assert mock_display.call_count == 0
+
+        # Test the "none" display method.
+        set_display(method="none")
+        assert SHOW_CONFIG["method"] == "none"
+        with patch("pygmt.figure.launch_external_viewer") as mock_viewer:
+            fig.show()
+            assert mock_viewer.call_count == 0
+        if _HAS_IPYTHON:
+            with patch("IPython.display.display") as mock_display:
+                fig.show()
+                assert mock_display.call_count == 0
 
         # Setting method to None should revert it to the default method.
         set_display(method=None)