test: expand comprehensive regression coverage

README.md

@@ -223,6 +223,7 @@ project does not promise bitwise GPU reproducibility.
 
 ```bash
 uv run pytest
+uv run pytest tests/unit
 GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity -v
 uv run ruff check .
 uv run ruff format --check .
@@ -235,6 +236,8 @@ and parity tests compare deterministic behavior with the frozen legacy
 reference. Performance guidance and measured caveats are in
 [`docs/performance.md`](docs/performance.md) and
 [`benchmarks/README.md`](benchmarks/README.md).
+See [`docs/testing.md`](docs/testing.md) for test layers, optional dependency
+markers, and package smoke validation.
 
 ## Architecture and migration status
 

docs/testing.md

@@ -0,0 +1,29 @@
+# Testing
+
+The suite is layered by dependency and purpose:
+
+```bash
+# Fast shared-package tests
+uv run pytest tests/unit
+
+# Full CPU suite, including integration and parity tests
+uv run pytest
+
+# Required ROOT-GNN parity gate (must fail if DGL is unavailable)
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest
+
+# Optional layers
+uv run pytest -m distributed -v
+uv run pytest -m onnx -v
+uv run pytest -m gpu -v
+GNN4COLLIDERS_ROOT_FIXTURE=/path/to/reduced.root uv run pytest -m real_data -v
+```
+
+Unit tests use deterministic, small tensors and generated ROOT files. DGL,
+ONNX, CUDA, distributed execution, and reduced real-data fixtures remain
+optional layers. Tests should assert public contracts and scientific
+invariants rather than private call sequences. New regression tests should
+use `tmp_path`, explicit seeds, and justified numerical tolerances.
+
+Release validation additionally builds a wheel and runs the CLI help command
+in a clean environment; it is not part of the ordinary pytest suite.

pyproject.toml

@@ -76,6 +76,10 @@ markers = [
     "distributed: tests requiring torch distributed execution",
     "legacy_env: tests requiring the historical runtime environment",
     "gpu: tests requiring CUDA",
+    "parity: comparisons with the frozen legacy implementation",
+    "onnx: tests requiring ONNX export/runtime dependencies",
+    "real_data: tests requiring an optional reduced ROOT fixture",
+    "slow: tests intentionally excluded from the fast development loop",
 ]
 
 [tool.ruff]

src/gnn4colliders/data/cache.py

@@ -38,7 +38,11 @@ class GraphSampleCache:
 
     def load(self) -> tuple[GraphSample, ...]:
         payload = torch.load(self.path, map_location="cpu", weights_only=False)
+        if not isinstance(payload, dict):
+            raise ValueError("cache payload must be a mapping")
         actual = payload.get("cache_metadata", {})
+        if not isinstance(actual, dict):
+            raise ValueError("cache metadata is missing or malformed")
         expected = asdict(self.metadata)
         for key in (
             "feature_schema_version",
@@ -50,7 +54,16 @@ class GraphSampleCache:
                     f"cache schema mismatch for {key}: "
                     f"found {actual.get(key)!r}, expected {expected[key]!r}"
                 )
-        return tuple(payload["samples"])
+        if actual.get("preprocessing", {}) != expected["preprocessing"]:
+            raise ValueError("cache preprocessing fingerprint mismatch")
+        if "samples" not in payload:
+            raise ValueError("cache samples are missing")
+        samples = payload["samples"]
+        if not isinstance(samples, (list, tuple)):
+            raise ValueError("cache samples are malformed")
+        if any(not isinstance(sample, GraphSample) for sample in samples):
+            raise ValueError("cache contains invalid graph samples")
+        return tuple(samples)
 
     def exists(self) -> bool:
         return self.path.is_file()

tests/integration/test_real_root_model.py

@@ -15,6 +15,8 @@ from gnn4colliders.graphs import build_dgl_graph
 from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
 
 pytest.importorskip("dgl")
+
+pytestmark = [pytest.mark.integration, pytest.mark.real_data]
 uproot = pytest.importorskip("uproot")
 
 FEATURE_BRANCHES = [

tests/integration/test_real_root_sample.py

@@ -15,6 +15,8 @@ from gnn4colliders.graphs import build_dgl_graph
 
 pytest.importorskip("dgl")
 
+pytestmark = [pytest.mark.integration, pytest.mark.real_data]
+
 FEATURE_BRANCHES = [
     ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"],
     ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0],

tests/parity/conftest.py

@@ -7,6 +7,10 @@ from pathlib import Path
 
 import pytest
 
+pytestmark = pytest.mark.parity
+
+pytestmark = pytest.mark.parity
+
 
 def _dgl_is_importable():
     try:

tests/unit/config/test_factories.py

@@ -0,0 +1,18 @@
+import pytest
+
+from gnn4colliders.config import build_task
+
+
+def test_factory_builds_semantic_task_names():
+    from gnn4colliders.tasks import (
+        BinaryClassificationTask,
+        MulticlassClassificationTask,
+    )
+
+    assert isinstance(build_task({"type": "binary"}), BinaryClassificationTask)
+    assert isinstance(build_task({"type": "multiclass"}), MulticlassClassificationTask)
+
+
+def test_factory_rejects_unknown_task_name():
+    with pytest.raises(ValueError, match="unsupported task"):
+        build_task({"type": "not-a-task"})

tests/unit/config/test_validation.py

@@ -0,0 +1,38 @@
+import pytest
+
+from gnn4colliders.config import validate_config
+
+
+@pytest.mark.parametrize(
+    ("section", "values", "message"),
+    [
+        ("data", {"batch_size": 0}, "batch_size"),
+        ("trainer", {"max_epochs": 0}, "max_epochs"),
+        ("model", {"out_size": 2}, "out_size"),
+    ],
+)
+def test_invalid_cross_config_values_fail_before_initialization(
+    section, values, message
+):
+    config = {"data": {}, "trainer": {}, "model": {}, "task": {"type": "binary"}}
+    config[section].update(values)
+    with pytest.raises(ValueError, match=message):
+        validate_config(config)
+
+
+def test_validation_rejects_conflicting_resume_and_transfer():
+    config = {"checkpoint": {"resume": "resume.pt", "pretrained": "base.pt"}}
+    with pytest.raises(ValueError, match="distinct workflows"):
+        validate_config(config)
+
+
+def test_validation_rejects_overlapping_split_folds():
+    config = {"data": {"splits": {"train_folds": [0], "test_folds": [0]}}}
+    with pytest.raises(ValueError, match="disjoint"):
+        validate_config(config)
+
+
+def test_fine_tuning_requires_pretrained_checkpoint():
+    config = {"model": {"name": "fine_tuned_edge_network"}}
+    with pytest.raises(ValueError, match="pretrained"):
+        validate_config(config)

tests/unit/data/test_batching.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
+from gnn4colliders.graphs import build_dgl_graph
+
+dgl = pytest.importorskip("dgl")
+
+
+def _sample(index: int, node_count: int) -> GraphSample:
+    nodes = torch.arange(node_count * 3, dtype=torch.float32).reshape(node_count, 3)
+    return GraphSample(
+        graph=build_dgl_graph(nodes),
+        label=torch.tensor(index),
+        global_features=torch.tensor([float(index), -float(index)]),
+        metadata=EventMetadata(index, index + 0.5, f"sample_{index}"),
+    )
+
+
+@pytest.mark.parametrize("counts", [(1,), (1, 2), (1, 3, 2)])
+def test_graph_batch_preserves_heterogeneous_order_and_metadata(counts):
+    samples = [_sample(index, count) for index, count in enumerate(counts)]
+    batch = batch_graph_samples(samples)
+    assert batch.labels.tolist() == list(range(len(counts)))
+    assert batch.metadata.sample_id == tuple(f"sample_{i}" for i in range(len(counts)))
+    assert batch.metadata.weight.tolist() == pytest.approx(
+        [i + 0.5 for i in range(len(counts))]
+    )
+    assert batch.global_features[:, 0].tolist() == list(map(float, range(len(counts))))
+    assert batch.graph.batch_num_nodes().tolist() == list(counts)
+    assert batch.graph.batch_num_edges().tolist() == [
+        count * (count - 1) if count > 1 else 1 for count in counts
+    ]

tests/unit/data/test_orchestration.py

@@ -5,12 +5,13 @@ import pytest
 
 from gnn4colliders.data import (
     FEATURE_SCHEMA_VERSION,
+    BatchMetadata,
     EventMetadata,
     SplitDefinition,
     select_folds,
 )
 from gnn4colliders.data.cache import CacheMetadata, GraphSampleCache
-from gnn4colliders.data.graph_dataset import GraphSample
+from gnn4colliders.data.graph_dataset import GraphDataLoader, GraphDataset, GraphSample
 
 
 def _sample(sample_id: str, fold: int) -> GraphSample:
@@ -34,6 +35,42 @@ def test_select_folds_is_named_deterministic_and_non_mutating():
     assert [sample.metadata.sample_id for sample in samples] == ["a", "b", "c"]
 
 
+def test_batch_metadata_preserves_order_and_extra_fields():
+    events = [
+        EventMetadata(2, 0.0, "sample_10", {"run": 10}),
+        EventMetadata(0, -3.5, "sample_2", {"run": 2, "tag": "b"}),
+    ]
+    batch = BatchMetadata.from_events(events)
+    assert batch.sample_id == ("sample_10", "sample_2")
+    assert batch.fold.tolist() == [2, 0]
+    assert batch.weight.tolist() == [0.0, -3.5]
+    assert batch.extra == {"run": (10, 2), "tag": (None, "b")}
+
+
+def test_graph_loader_shuffle_is_seeded_by_epoch_without_global_rng_mutation():
+    pytest.importorskip("dgl")
+    import torch
+
+    from gnn4colliders.graphs import build_dgl_graph
+
+    samples = [
+        GraphSample(
+            build_dgl_graph(torch.ones(1, 3)),
+            torch.tensor(i),
+            None,
+            EventMetadata(0, 1.0, str(i)),
+        )
+        for i in range(5)
+    ]
+    loader = GraphDataLoader(GraphDataset(samples), 2, shuffle=True, seed=11)
+    first = [item for batch in loader for item in batch.metadata.sample_id]
+    second = [item for batch in loader for item in batch.metadata.sample_id]
+    loader.set_epoch(1)
+    third = [item for batch in loader for item in batch.metadata.sample_id]
+    assert first == second
+    assert third != first
+
+
 def test_split_definition_rejects_overlapping_folds():
     with pytest.raises(ValueError, match="disjoint"):
         SplitDefinition(train_folds=frozenset({0}), test_folds=frozenset({0}))
@@ -53,3 +90,45 @@ def test_graph_cache_round_trip_and_schema_check(tmp_path: Path):
     )
     with pytest.raises(ValueError, match="schema mismatch"):
         stale.load()
+
+
+@pytest.mark.parametrize(
+    "payload, message",
+    [
+        ({"samples": ()}, "schema mismatch"),
+        ({"cache_metadata": {}, "samples": ()}, "schema mismatch"),
+        (
+            {
+                "cache_metadata": {
+                    "feature_schema_version": 1,
+                    "graph_schema_version": 1,
+                    "cache_schema_version": 1,
+                }
+            },
+            "samples",
+        ),
+    ],
+)
+def test_cache_rejects_missing_or_incomplete_metadata(tmp_path, payload, message):
+    path = tmp_path / "broken.pt"
+    import torch
+
+    torch.save(payload, path)
+    with pytest.raises(ValueError, match=message):
+        GraphSampleCache(path).load()
+
+
+def test_cache_rejects_preprocessing_fingerprint_mismatch(tmp_path):
+    path = tmp_path / "graphs.pt"
+    GraphSampleCache(path, CacheMetadata(preprocessing={"features": "a"})).save(
+        [_sample("event:0", 0)]
+    )
+    with pytest.raises(ValueError, match="fingerprint"):
+        GraphSampleCache(path, CacheMetadata(preprocessing={"features": "b"})).load()
+
+
+def test_cache_rejects_truncated_artifact(tmp_path):
+    path = tmp_path / "truncated.pt"
+    path.write_bytes(b"not a torch archive")
+    with pytest.raises(Exception):
+        GraphSampleCache(path).load()

tests/unit/data/test_root_dataset.py

@@ -89,3 +89,37 @@ def test_multiple_files_keep_input_order_and_feed_feature_builder(tmp_path: Path
     )
     assert features.shape == (2, 4)
     assert lengths == [2]
+
+
+def test_root_boundary_reports_missing_tree_and_branch(tmp_path: Path):
+    path = tmp_path / "events.root"
+    _write_fixture(path)
+    with pytest.raises(KeyError, match="tree"):
+        read_tree(path, "missing", ["weight"])
+    with pytest.raises(KeyError, match="missing requested branches"):
+        read_tree(path, "events", ["missing_branch"])
+
+
+def test_root_boundary_reports_missing_file(tmp_path: Path):
+    with pytest.raises((OSError, FileNotFoundError)):
+        read_tree(tmp_path / "does-not-exist.root", "events", ["weight"])
+
+
+def test_sample_identity_is_stable_and_distinguishes_files(tmp_path: Path):
+    first = tmp_path / "first.root"
+    second = tmp_path / "second.root"
+    _write_fixture(first)
+    _write_fixture(second)
+    kwargs = dict(
+        tree_name="events",
+        label="label_branch",
+        feature_branches=[["jet_pt"]],
+        fold_var="eventNumber",
+        weight_var="weight",
+    )
+    left = RootEventDataset(first, **kwargs)
+    repeated = RootEventDataset(first, **kwargs)
+    both = RootEventDataset([first, second], **kwargs)
+    assert left[0].metadata.sample_id == repeated[0].metadata.sample_id
+    assert left[0].metadata.sample_id != left[1].metadata.sample_id
+    assert both[0].metadata.sample_id != both[3].metadata.sample_id

tests/unit/distributed/test_ddp_cpu.py

@@ -8,6 +8,8 @@ from torch import nn
 
 from gnn4colliders.distributed import finalize, initialize, prepare_model
 
+pytestmark = pytest.mark.distributed
+
 
 def _port() -> int:
     with socket.socket() as sock:

tests/unit/export/test_onnx_export.py

@@ -1,8 +1,6 @@
 import pytest
 import torch
 
-pytest.importorskip("dgl")
-
 from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples
 from gnn4colliders.export import (
     RootGNNExportAdapter,
@@ -12,6 +10,9 @@ from gnn4colliders.export import (
 from gnn4colliders.graphs import build_dgl_graph
 from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
 
+pytestmark = pytest.mark.onnx
+dgl = pytest.importorskip("dgl")
+
 
 def _batch():
     samples = []

tests/unit/features/test_objects.py

@@ -87,3 +87,24 @@ def test_empty_vectors_and_input_immutability(event, schema):
     assert lengths == [0, 1, 1, 1, 1]
     assert features.shape == (4, 7)
     np.testing.assert_array_equal(event["ele_pt"], before)
+
+
+def test_reordering_vector_objects_reorders_feature_rows(event, schema):
+    names, object_types, scales = schema
+    reordered = dict(event)
+    for name in ("jet_pt", "jet_eta", "jet_phi", "jet_btag"):
+        reordered[name] = event[name][::-1].copy()
+    first, _ = build_node_features(event, names, object_types, scales)
+    second, _ = build_node_features(reordered, names, object_types, scales)
+    torch.testing.assert_close(second[:2], first[[1, 0]])
+    torch.testing.assert_close(second[2:], first[2:])
+
+
+def test_feature_scaling_changes_only_the_requested_columns(event, schema):
+    names, object_types, scales = schema
+    base, _ = build_node_features(event, names, object_types, scales)
+    changed_scales = list(scales)
+    changed_scales[0] *= 2
+    scaled, _ = build_node_features(event, names, object_types, changed_scales)
+    torch.testing.assert_close(scaled[:, 0], base[:, 0] * 2)
+    torch.testing.assert_close(scaled[:, 1:], base[:, 1:])

tests/unit/graphs/test_edge_features.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from gnn4colliders.graphs import build_edge_features, fully_connected_edges
@@ -18,3 +19,33 @@ def test_empty_edges_keep_three_feature_columns():
     features = build_edge_features(nodes, src, dst, eta_index=1, phi_index=2)
     assert features.shape == (0, 3)
     assert features.dtype == torch.float32
+
+
+def test_edge_features_have_symmetry_and_nonnegative_distance():
+    nodes = torch.tensor([[1.0, 2.0, 3.13], [2.0, -1.0, -3.13], [3.0, 0.5, 0.2]])
+    src, dst = fully_connected_edges(3)
+    actual = build_edge_features(nodes, src, dst, eta_index=1, phi_index=2)
+    reverse = build_edge_features(nodes, dst, src, eta_index=1, phi_index=2)
+    torch.testing.assert_close(actual[:, 0], -reverse[:, 0])
+    torch.testing.assert_close(actual[:, 2], reverse[:, 2])
+    assert torch.all(actual[:, 2] >= 0)
+    assert torch.all(actual[:, 1].abs() <= torch.pi)
+
+
+@pytest.mark.parametrize("bad_nodes", [torch.zeros(3), torch.zeros(2, 2)])
+def test_edge_features_reject_malformed_node_features(bad_nodes):
+    with pytest.raises((ValueError, IndexError)):
+        build_edge_features(
+            bad_nodes, torch.tensor([0]), torch.tensor([0]), eta_index=1, phi_index=2
+        )
+
+
+def test_edge_features_reject_mismatched_edge_indices():
+    with pytest.raises(ValueError, match="equal shape"):
+        build_edge_features(
+            torch.zeros(2, 3),
+            torch.tensor([0]),
+            torch.tensor([1, 0]),
+            eta_index=1,
+            phi_index=2,
+        )

tests/unit/graphs/test_topology.py

@@ -50,3 +50,29 @@ def test_self_loops_can_be_enabled_and_zero_nodes_is_empty():
 def test_negative_node_count_is_rejected():
     with pytest.raises(ValueError, match="non-negative"):
         fully_connected_edges(-1)
+
+
+@pytest.mark.parametrize("num_nodes", [0, 1, 2, 4, 7])
+def test_topology_invariants_hold_for_multiple_graph_sizes(num_nodes):
+    src, dst = fully_connected_edges(num_nodes)
+    expected = (
+        0 if num_nodes == 0 else 1 if num_nodes == 1 else num_nodes * (num_nodes - 1)
+    )
+    assert src.numel() == dst.numel() == expected
+    pairs = list(zip(src.tolist(), dst.tolist()))
+    assert len(pairs) == len(set(pairs))
+    if num_nodes > 1:
+        assert all(source != destination for source, destination in pairs)
+        assert [src.tolist().count(index) for index in range(num_nodes)] == [
+            num_nodes - 1
+        ] * num_nodes
+        assert [dst.tolist().count(index) for index in range(num_nodes)] == [
+            num_nodes - 1
+        ] * num_nodes
+
+
+def test_topology_rejects_bool_and_non_integer_counts():
+    with pytest.raises(TypeError, match="integer"):
+        fully_connected_edges(True)
+    with pytest.raises(TypeError, match="integer"):
+        fully_connected_edges(2.0)

tests/unit/inference/test_predictor.py

@@ -81,3 +81,27 @@ def test_npz_round_trip(tmp_path):
         np.testing.assert_array_equal(values["logits"], result.logits.numpy())
         np.testing.assert_array_equal(values["labels"], result.labels.numpy())
         np.testing.assert_array_equal(values["fold"], result.fold.numpy())
+
+
+def test_predictor_is_reentrant_and_returns_detached_cpu_results():
+    predictor = Predictor(_Model([[-1.0], [1.0]]), BinaryClassificationTask())
+    first = predictor.predict(_loader([[-1.0], [1.0]], [0, 1]))
+    second = predictor.predict(_loader([[-1.0], [1.0]], [0, 1]))
+    torch.testing.assert_close(first.logits, second.logits)
+    assert first.logits.device.type == "cpu"
+    assert not first.logits.requires_grad
+    assert first.logits.grad_fn is None
+
+
+def test_npz_omits_labels_for_unlabeled_results(tmp_path):
+    metadata = BatchMetadata(
+        fold=torch.tensor([0]), weight=torch.tensor([1.0]), sample_id=("sample_7",)
+    )
+    batch = SimpleNamespace(metadata=metadata, offset=0, to=lambda device: batch)
+    result = Predictor(_Model([[2.0, 1.0]]), MulticlassClassificationTask()).predict(
+        [batch]
+    )
+    path = write_npz(result, tmp_path / "unlabeled.npz")
+    with np.load(path) as values:
+        assert "labels" not in values
+        assert values["sample_id"].tolist() == ["sample_7"]

tests/unit/models/root_gnn/test_edge_network.py

@@ -39,3 +39,40 @@ def test_edge_network_supports_multiclass_and_transfer_freezing():
         parameter.requires_grad for parameter in transferred.classifier.parameters()
     )
     assert model(graph, torch.zeros(1, 2)).shape == (1, 12)
+
+
+@pytest.mark.parametrize("out_size", [1, 3])
+@pytest.mark.parametrize("n_proc_steps", [0, 1])
+def test_model_shape_matrix_and_state_dict_round_trip(out_size, n_proc_steps):
+    graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]))
+    globals_ = torch.tensor([[1.0, 2.0]])
+    torch.manual_seed(7)
+    model = EdgeNetwork(graph, globals_, 4, out_size, 1, n_proc_steps).eval()
+    with torch.no_grad():
+        expected = model(graph, globals_)
+    fresh = EdgeNetwork(graph, globals_, 4, out_size, 1, n_proc_steps).eval()
+    fresh.load_state_dict(model.state_dict())
+    with torch.no_grad():
+        actual = fresh(graph, globals_)
+    torch.testing.assert_close(actual, expected)
+
+
+def test_model_backward_produces_finite_gradients_and_update():
+    model, graph = _model()
+    model.train()
+    before = [parameter.detach().clone() for parameter in model.parameters()]
+    loss = model(graph, torch.zeros(1, 2)).square().mean()
+    assert torch.isfinite(loss)
+    loss.backward()
+    gradients = [
+        parameter.grad for parameter in model.parameters() if parameter.requires_grad
+    ]
+    assert gradients and all(
+        gradient is not None and torch.isfinite(gradient).all()
+        for gradient in gradients
+    )
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    optimizer.step()
+    assert any(
+        not torch.equal(old, new) for old, new in zip(before, model.parameters())
+    )

tests/unit/tasks/test_classification.py

@@ -64,3 +64,21 @@ def test_multiclass_loss_and_metrics():
     output = task.predict(logits)
     assert output["predictions"].tolist() == [0, 1, 2]
     assert task.metrics(logits, batch)["accuracy"] == 1.0
+
+
+def test_binary_loss_is_invariant_under_joint_event_permutation():
+    task = BinaryClassificationTask()
+    batch = _batch([0, 1, 0, 1], [1.0, 2.0, 3.0, 4.0])
+    logits = torch.tensor([[-2.0], [1.0], [0.5], [-0.5]])
+    permutation = torch.tensor([2, 0, 3, 1])
+    shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist())
+    assert task.loss(logits, batch) == pytest.approx(task.loss(logits[permutation], shuffled).item())
+
+
+def test_multiclass_metrics_are_invariant_under_joint_event_permutation():
+    task = MulticlassClassificationTask()
+    batch = _batch([0, 1, 2, 1], [1.0, 2.0, 1.0, 3.0])
+    logits = torch.tensor([[3.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 4.0], [0.0, 1.0, 0.0]])
+    permutation = torch.tensor([3, 0, 2, 1])
+    shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist())
+    assert task.metrics(logits, batch) == task.metrics(logits[permutation], shuffled)

tests/unit/training/test_checkpoint.py

@@ -1,6 +1,7 @@
 import random
 
 import numpy as np
+import pytest
 import torch
 from torch import nn
 
@@ -105,3 +106,22 @@ def test_rng_state_is_restored(tmp_path):
     assert expected[0] == actual[0]
     assert expected[1] == actual[1]
     assert torch.equal(expected[2], actual[2])
+
+
+def test_checkpoint_rejects_unsupported_schema_and_missing_weights(tmp_path):
+    unsupported = tmp_path / "unsupported.pt"
+    torch.save({"schema_version": 999}, unsupported)
+    with pytest.raises(ValueError, match="unsupported checkpoint schema"):
+        CheckpointManager.load(unsupported)
+
+    missing = tmp_path / "missing.pt"
+    torch.save({"schema_version": 1, "epoch": 0}, missing)
+    with pytest.raises(ValueError, match="incompatible"):
+        load_model_weights(nn.Linear(2, 1), CheckpointManager.load(missing))
+
+
+def test_checkpoint_loading_rejects_architecture_mismatch(tmp_path):
+    manager = CheckpointManager(tmp_path)
+    path = manager.save(model=nn.Linear(2, 1), trainer_state=TrainerState(epoch=0))
+    with pytest.raises(ValueError, match="incompatible"):
+        load_model_weights(nn.Linear(3, 1), manager.load(path))

tests/unit/training/test_trainer.py

@@ -1,3 +1,4 @@
+import pytest
 import torch
 from torch import nn
 
@@ -79,3 +80,13 @@ def test_fit_uses_validation_for_early_stopping_and_records_history():
     assert len(history.train) == len(history.validation)
     assert 1 <= len(history.train) <= 4
     assert all(isinstance(item.loss, float) for item in history.train)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable")
+def test_graph_batch_to_cuda_preserves_named_metadata():
+    batch = _batch([1], [0])
+    moved = batch.to("cuda")
+    assert moved.labels.device.type == "cuda"
+    assert moved.metadata.weight.device.type == "cuda"
+    assert moved.metadata.sample_id == ("event:0",)