Split test suite into unit/ and integration/ subdirectories

.github/workflows/install-and-test.yml

@@ -114,10 +114,20 @@ jobs:
         run: |
           python -c "import neural_lam; print('neural-lam imported successfully')"
 
-      - name: Run tests
+      - name: Run tests (unit)
         if: matrix.package_manager == 'uv'
         run: |
-          pytest -vv -s --doctest-modules
+          pytest -vv -s tests/unit/ --tb=short
+
+      - name: Run tests (integration)
+        if: matrix.package_manager == 'uv'
+        run: |
+          pytest -vv -s tests/integration/ --tb=short
+
+      - name: Run doctests
+        if: matrix.package_manager == 'uv'
+        run: |
+          pytest -vv --doctest-modules neural_lam/
 
       - name: Upload test figures
         if: matrix.package_manager == 'uv'

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased](https://github.com/mllam/neural-lam/compare/v0.6.0...HEAD)
 
+### Maintenance
+
+- Split test suite into `tests/unit/` and `tests/integration/`, moving the S3 fixture into the integration conftest and refactoring `test_clamping.py` to use `DummyDatastore` [\#615](https://github.com/mllam/neural-lam/pull/615) @jishanahmed-shaikh
+
 ### Fixed
 
 - Fix `AssertionError` in `aggregate_and_plot_metrics` when using `--metrics_watch` flags by using `isinstance` dispatch for figure vs scalar logging [\#303](https://github.com/mllam/neural-lam/pull/303) @AftAb-25

pyproject.toml

@@ -108,6 +108,9 @@ allow-any-import-level = "neural_lam"
 [tool.pylint.SIMILARITIES]
 min-similarity-lines = 10
 
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
 [build-system]
 requires = ["hatchling>=1.27.0", "hatch-vcs"]
 build-backend = "hatchling.build"

tests/conftest.py

@@ -1,19 +1,12 @@
 # Standard library
 import os
-from datetime import timedelta
-from pathlib import Path
 
 # Third-party
-import pooch
 import pytest
-import yaml
 from pytorch_lightning.utilities import rank_zero_only
 
 # First-party
-from neural_lam.datastore import DATASTORES, init_datastore
-from neural_lam.datastore.npyfilesmeps import (
-    compute_standardization_stats as compute_standardization_stats_meps,
-)
+from neural_lam.datastore import DATASTORES
 
 # Local
 from .dummy_datastore import DummyDatastore
@@ -22,104 +15,13 @@
 # and to avoid having to deal with authentication
 os.environ["WANDB_MODE"] = "disabled"
 
+# Register DummyDatastore so DatastoreSelection validation passes
+# in any test that uses it (unit or integration)
+DATASTORES.setdefault(DummyDatastore.SHORT_NAME, DummyDatastore)
+
 
 @pytest.fixture(autouse=True)
 def ensure_rank_zero(monkeypatch):
     """Ensure rank_zero_only.rank == 0 so @rank_zero_only-decorated functions
     execute their body regardless of state left by prior training tests."""
     monkeypatch.setattr(rank_zero_only, "rank", 0, raising=False)
-
-
-DATASTORE_EXAMPLES_ROOT_PATH = Path("tests/datastore_examples")
-
-# Initializing variables for the s3 client
-S3_BUCKET_NAME = "mllam-testdata"
-S3_ENDPOINT_URL = "https://object-store.os-api.cci1.ecmwf.int"
-S3_FILE_PATH = "neural-lam/npy/meps_example_reduced.v0.3.0.tar.gz"
-S3_FULL_PATH = "/".join([S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_FILE_PATH])
-TEST_DATA_KNOWN_HASH = (
-    "af16d87099944dda152cc988ce347173da68c56d3739c86ec53c8132981a93e6"
-)
-
-
-def download_meps_example_reduced_dataset():
-    # Download and unzip test data into data/meps_example_reduced
-    root_path = DATASTORE_EXAMPLES_ROOT_PATH / "npyfilesmeps"
-    dataset_path = root_path / "meps_example_reduced"
-
-    pooch.retrieve(
-        url=S3_FULL_PATH,
-        known_hash=TEST_DATA_KNOWN_HASH,
-        processor=pooch.Untar(extract_dir=""),
-        path=root_path,
-        fname=S3_FILE_PATH.split("/")[-1],
-    )
-
-    config_path = dataset_path / "meps_example_reduced.datastore.yaml"
-
-    with open(config_path, "r") as f:
-        config = yaml.safe_load(f)
-
-    if "class" in config["projection"]:
-        # XXX: should update the dataset stored on S3 with the change below
-        #
-        # rename the `projection.class` key to `projection.class_name` in the
-        # config this is because the `class` key is reserved for the class
-        # attribute of the object and so we can't use it to define a python
-        # dataclass
-        config["projection"]["class_name"] = config["projection"].pop("class")
-
-        with open(config_path, "w") as f:
-            yaml.dump(config, f)
-
-    # create parameters, only run if the files we expect are not present
-    expected_parameter_files = [
-        "parameter_mean.pt",
-        "parameter_std.pt",
-        "diff_mean.pt",
-        "diff_std.pt",
-    ]
-    expected_parameter_filepaths = [
-        dataset_path / "static" / fn for fn in expected_parameter_files
-    ]
-    if any(not p.exists() for p in expected_parameter_filepaths):
-        compute_standardization_stats_meps.main(
-            datastore_config_path=config_path,
-            batch_size=8,
-            step_length=timedelta(hours=3),
-            n_workers=0,
-            distributed=False,
-        )
-
-    return config_path
-
-
-DATASTORES_EXAMPLES = dict(
-    mdp=(
-        DATASTORE_EXAMPLES_ROOT_PATH
-        / "mdp"
-        / "danra_100m_winds"
-        / "danra.datastore.yaml"
-    ),
-    npyfilesmeps=None,
-    dummydata=None,
-)
-
-DATASTORES[DummyDatastore.SHORT_NAME] = DummyDatastore
-
-
-def init_datastore_example(datastore_kind):
-    if (
-        datastore_kind == "npyfilesmeps"
-        and DATASTORES_EXAMPLES["npyfilesmeps"] is None
-    ):
-        DATASTORES_EXAMPLES["npyfilesmeps"] = (
-            download_meps_example_reduced_dataset()
-        )
-
-    datastore = init_datastore(
-        datastore_kind=datastore_kind,
-        config_path=DATASTORES_EXAMPLES[datastore_kind],
-    )
-
-    return datastore

tests/integration/conftest.py

@@ -0,0 +1,110 @@
+# Standard library
+from datetime import timedelta
+from pathlib import Path
+
+# Third-party
+import pooch
+import yaml
+
+# First-party
+from neural_lam.datastore import DATASTORES, init_datastore
+from neural_lam.datastore.npyfilesmeps import (
+    compute_standardization_stats as compute_standardization_stats_meps,
+)
+
+# Local
+from tests.dummy_datastore import DummyDatastore
+
+DATASTORE_EXAMPLES_ROOT_PATH = Path("tests/datastore_examples")
+
+# Initializing variables for the s3 client
+S3_BUCKET_NAME = "mllam-testdata"
+S3_ENDPOINT_URL = "https://object-store.os-api.cci1.ecmwf.int"
+S3_FILE_PATH = "neural-lam/npy/meps_example_reduced.v0.3.0.tar.gz"
+S3_FULL_PATH = "/".join([S3_ENDPOINT_URL, S3_BUCKET_NAME, S3_FILE_PATH])
+TEST_DATA_KNOWN_HASH = (
+    "af16d87099944dda152cc988ce347173da68c56d3739c86ec53c8132981a93e6"
+)
+
+
+def download_meps_example_reduced_dataset():
+    # Download and unzip test data into data/meps_example_reduced
+    root_path = DATASTORE_EXAMPLES_ROOT_PATH / "npyfilesmeps"
+    dataset_path = root_path / "meps_example_reduced"
+
+    pooch.retrieve(
+        url=S3_FULL_PATH,
+        known_hash=TEST_DATA_KNOWN_HASH,
+        processor=pooch.Untar(extract_dir=""),
+        path=root_path,
+        fname=S3_FILE_PATH.split("/")[-1],
+    )
+
+    config_path = dataset_path / "meps_example_reduced.datastore.yaml"
+
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    if "class" in config["projection"]:
+        # XXX: should update the dataset stored on S3 with the change below
+        #
+        # rename the `projection.class` key to `projection.class_name` in the
+        # config this is because the `class` key is reserved for the class
+        # attribute of the object and so we can't use it to define a python
+        # dataclass
+        config["projection"]["class_name"] = config["projection"].pop("class")
+
+        with open(config_path, "w") as f:
+            yaml.dump(config, f)
+
+    # create parameters, only run if the files we expect are not present
+    expected_parameter_files = [
+        "parameter_mean.pt",
+        "parameter_std.pt",
+        "diff_mean.pt",
+        "diff_std.pt",
+    ]
+    expected_parameter_filepaths = [
+        dataset_path / "static" / fn for fn in expected_parameter_files
+    ]
+    if any(not p.exists() for p in expected_parameter_filepaths):
+        compute_standardization_stats_meps.main(
+            datastore_config_path=config_path,
+            batch_size=8,
+            step_length=timedelta(hours=3),
+            n_workers=0,
+            distributed=False,
+        )
+
+    return config_path
+
+
+DATASTORES_EXAMPLES = dict(
+    mdp=(
+        DATASTORE_EXAMPLES_ROOT_PATH
+        / "mdp"
+        / "danra_100m_winds"
+        / "danra.datastore.yaml"
+    ),
+    npyfilesmeps=None,
+    dummydata=None,
+)
+
+DATASTORES[DummyDatastore.SHORT_NAME] = DummyDatastore
+
+
+def init_datastore_example(datastore_kind):
+    if (
+        datastore_kind == "npyfilesmeps"
+        and DATASTORES_EXAMPLES["npyfilesmeps"] is None
+    ):
+        DATASTORES_EXAMPLES["npyfilesmeps"] = (
+            download_meps_example_reduced_dataset()
+        )
+
+    datastore = init_datastore(
+        datastore_kind=datastore_kind,
+        config_path=DATASTORES_EXAMPLES[datastore_kind],
+    )
+
+    return datastore

tests/test_datasets.py → tests/integration/test_datasets.py

@@ -14,7 +14,7 @@
 from neural_lam.datastore.base import BaseRegularGridDatastore
 from neural_lam.models.graph_lam import GraphLAM
 from neural_lam.weather_dataset import WeatherDataset
-from tests.conftest import init_datastore_example
+from tests.integration.conftest import init_datastore_example
 from tests.dummy_datastore import DummyDatastore, EnsembleDummyDatastore
 
 

tests/test_datastores.py → tests/integration/test_datastores.py

@@ -50,7 +50,7 @@
 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
+from tests.integration.conftest import init_datastore_example
 
 
 @pytest.mark.parametrize("datastore_name", DATASTORES.keys())

tests/test_graph_creation.py → tests/integration/test_graph_creation.py

@@ -10,7 +10,7 @@
 from neural_lam.create_graph import create_graph_from_datastore
 from neural_lam.datastore import DATASTORES
 from neural_lam.datastore.base import BaseRegularGridDatastore
-from tests.conftest import init_datastore_example
+from tests.integration.conftest import init_datastore_example
 
 
 @pytest.mark.parametrize("graph_name", ["1level", "multiscale", "hierarchical"])

tests/test_plotting.py → tests/integration/test_plotting.py

@@ -19,7 +19,7 @@
 from neural_lam.create_graph import create_graph_from_datastore
 from neural_lam.models.graph_lam import GraphLAM
 from neural_lam.weather_dataset import WeatherDataset
-from tests.conftest import init_datastore_example
+from tests.integration.conftest import init_datastore_example
 from tests.dummy_datastore import DummyDatastore
 
 # Create output directory for test figures

tests/test_training.py → tests/integration/test_training.py

@@ -16,7 +16,7 @@
 from neural_lam.models.ar_model import ARModel
 from neural_lam.models.graph_lam import GraphLAM
 from neural_lam.weather_dataset import WeatherDataModule
-from tests.conftest import init_datastore_example
+from tests.integration.conftest import init_datastore_example
 
 
 def run_simple_training(datastore, set_output_std, metrics_watch=None):