fix: tests

ecmwf · Jun 27, 2024 · b28b0ff · b28b0ff
1 parent fcd0b74
commit b28b0ff
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/nodes/test_npz.py b/tests/nodes/test_npz.py
@@ -41,7 +41,7 @@ def test_register_nodes(mock_grids_path: str, resolution: str):
     assert graph["test_nodes"].x is not None
     assert isinstance(graph["test_nodes"].x, torch.Tensor)
     assert graph["test_nodes"].x.shape == (num_nodes, 2)
-    assert graph["test_nodes"].node_type == "NPZNodes"
+    assert graph["test_nodes"].node_type == "NPZFileNodeBuilder"
 
 
 @pytest.mark.parametrize("attr_class", [UniformWeights, AreaWeights])

diff --git a/tests/nodes/test_zarr.py b/tests/nodes/test_zarr.py
@@ -33,7 +33,7 @@ def test_register_nodes(mocker, mock_zarr_dataset):
     assert graph["test_nodes"].x is not None
     assert isinstance(graph["test_nodes"].x, torch.Tensor)
     assert graph["test_nodes"].x.shape == (node_builder.ds.num_nodes, 2)
-    assert graph["test_nodes"].node_type == "ZarrNodes"
+    assert graph["test_nodes"].node_type == "ZarrDatasetNodeBuilder"
 
 
 @pytest.mark.parametrize("attr_class", [UniformWeights, AreaWeights])