From 89b10b58d035e90f95997a57ea9b0cd27e55d70e Mon Sep 17 00:00:00 2001
From: Leif Denby <leif@denby.eu>
Date: Tue, 12 Nov 2024 20:44:23 +0100
Subject: [PATCH] add test for datastores example plot function

---
 tests/test_datastores.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tests/test_datastores.py b/tests/test_datastores.py
index 00cd508..c28418c 100644
--- a/tests/test_datastores.py
+++ b/tests/test_datastores.py
@@ -48,6 +48,7 @@
 # First-party
 from neural_lam.datastore import DATASTORES
 from neural_lam.datastore.base import BaseRegularGridDatastore
+from neural_lam.datastore.plot_example import plot_example_from_datastore
 from tests.conftest import init_datastore_example
 
 
@@ -364,3 +365,21 @@ def test_dataarray_shapes(datastore_name):
     )
     # assert diff == 0 with tolerance 1e-6
     assert torch.allclose(diff, torch.zeros_like(diff), atol=1e-6)
+
+
+@pytest.mark.parametrize("datastore_name", DATASTORES.keys())
+def test_plot_example_from_datastore(datastore_name):
+    """Check that the `plot_example_from_datastore` function is implemented."""
+    datastore = init_datastore_example(datastore_name)
+    fig = plot_example_from_datastore(
+        category="static",
+        datastore=datastore,
+        col_dim="{category}_feature",
+        split="train",
+        standardize=True,
+        selection={},
+        index_selection={},
+    )
+
+    assert fig is not None
+    assert fig.get_axes()