Improve the tests

Author: seisman

pygmt/clib/session.py

@@ -1690,7 +1690,7 @@ def virtualfile_from_stringio(self, stringio: io.StringIO):
         for i, segment in enumerate(segments):
             seg = table.segment[i].contents
             if segment["header"] != "":
-                seg.header = segment["header"].encode() # type: ignore[attr-defined]
+                seg.header = segment["header"].encode()  # type: ignore[attr-defined]
             seg.text = strings_to_ctypes_array(segment["data"])
 
         with self.open_virtualfile(family, geometry, "GMT_IN", dataset) as vfile:

pygmt/tests/test_clib_virtualfiles.py

@@ -415,62 +415,99 @@ class TestVirtualfileFromStringIO:
     Test the virtualfile_from_stringio method.
     """
 
-    def _check_virtualfile_from_stringio(self, data: str):
+    def _stringio_to_dataset(self, data: io.StringIO):
         """
-        A helper function to check the output of the virtualfile_from_stringio method.
+        A helper function for check the virtualfile_from_stringio method.
+
+        The function does the following:
+
+        1. Creates a virtual file from the input StringIO object.
+        2. Pass the virtual file to the ``read`` module, which reads the virtual file
+           and writes it to another virtual file.
+        3. Reads the output virtual file as a GMT_DATASET object.
+        4. Extracts the header and the trailing text from the dataset and returns it as
+           a string.
         """
-        # The expected output is the data with all comment lines removed.
-        expected = (
-            "\n".join(line for line in data.splitlines() if not line.startswith("#"))
-            + "\n"
-        )
-        stringio = io.StringIO(data)
-        with GMTTempFile() as outfile:
-            with clib.Session() as lib:
-                with lib.virtualfile_from_stringio(stringio) as vintbl:
-                    lib.call_module("write", args=[vintbl, f"->{outfile.name}", "-Td"])
-                    output = outfile.read()
-        assert output == expected
+        with clib.Session() as lib:
+            with (
+                lib.virtualfile_from_stringio(data) as vintbl,
+                lib.virtualfile_out(kind="dataset") as vouttbl,
+            ):
+                lib.call_module("read", args=[vintbl, vouttbl, "-Td"])
+                ds = lib.read_virtualfile(vouttbl, kind="dataset").contents
+
+                output = []
+                table = ds.table[0].contents
+                for segment in table.segment[: table.n_segments]:
+                    seg = segment.contents
+                    output.append(f"> {seg.header.decode()}" if seg.header else ">")
+                    output.extend(np.char.decode(seg.text[: seg.n_rows]))
+            return "\n".join(output) + "\n"
 
     def test_virtualfile_from_stringio(self):
         """
         Test the virtualfile_from_stringio method.
         """
-        data = (
+        data = io.StringIO(
             "# Comment\n"
             "H 24p Legend\n"
             "N 2\n"
             "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
         )
-        self._check_virtualfile_from_stringio(data)
+        expected = (
+            ">\n"
+            "H 24p Legend\n"
+            "N 2\n"
+            "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+        )
+        assert self._stringio_to_dataset(data) == expected
 
     def test_one_segment(self):
         """
         Test the virtualfile_from_stringio method with one segment.
         """
-        data = (
+        data = io.StringIO(
             "# Comment\n"
             "> Segment 1\n"
-            "H 24p Legend\n"
-            "N 2\n"
-            "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FGHIJK LMN OPQ\n"
+            "RSTUVWXYZ\n"
         )
-        self._check_virtualfile_from_stringio(data)
+        expected = (
+            "> Segment 1\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FGHIJK LMN OPQ\n"
+            "RSTUVWXYZ\n"
+        )
+        assert self._stringio_to_dataset(data) == expected
 
     def test_multiple_segments(self):
         """
         Test the virtualfile_from_stringio method with multiple segments.
         """
-        data = (
+        data = io.StringIO(
             "# Comment line 1\n"
             "# Comment line 2\n"
             "> Segment 1\n"
-            "H 24p Legend\n"
-            "N 2\n"
-            "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FG\n"
+            "# Comment line 3\n"
             "> Segment 2\n"
-            "H 24p Legend\n"
-            "N 2\n"
-            "S 0.1i c 0.15i p300/12 0.25p 0.3i My circle\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FG\n"
+        )
+        expected = (
+            "> Segment 1\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FG\n"
+            "> Segment 2\n"
+            "1 2 3 ABC\n"
+            "4 5 DE\n"
+            "6 7 8   9  FG\n"
         )
-        self._check_virtualfile_from_stringio(data)
+        assert self._stringio_to_dataset(data) == expected