Add test coverage for to_pickle export

CHANGELOG.md

@@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add `__version__` attribute to the package init
   [\#56](https://github.com/mllam/weather-model-graphs/pull/56) @AdMub
 
+- Add test coverage for `to_pickle` export in `tests/test_save.py`
+
 ## [v0.3.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.3.0)
 
 

tests/test_save.py

@@ -1,4 +1,8 @@
+import pickle
 import tempfile
+from pathlib import Path
+
+import networkx as nx
 
 import pytest
 from loguru import logger
@@ -40,3 +44,26 @@ def test_save_to_pyg(list_from_attribute):
                 name=name,
                 list_from_attribute=list_from_attribute,
             )
+
+
+def test_save_to_pickle():
+    """
+    Test that to_pickle safely exports a graph to disk and allows
+    reloading without structural loss.
+    """
+    xy = test_utils.create_fake_xy(N=64)
+    graph = wmg.create.archetype.create_oskarsson_hierarchical_graph(coords=xy)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        name = "test_pickle_graph"
+        wmg.save.to_pickle(graph=graph, output_directory=tmpdir, name=name)
+
+        expected_path = Path(tmpdir) / f"{name}.pickle"
+        assert expected_path.exists()
+
+        with open(expected_path, "rb") as f:
+            loaded_graph = pickle.load(f)
+
+        assert isinstance(loaded_graph, nx.DiGraph)
+        assert len(loaded_graph.nodes) == len(graph.nodes)
+        assert len(loaded_graph.edges) == len(graph.edges)