diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cf3ab0..9f18160 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/tests/test_save.py b/tests/test_save.py index f889ab9..684165a 100644 --- a/tests/test_save.py +++ b/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)