fixing pytest

ufuk-cakir · Nov 6, 2024 · 9f472b2 · 9f472b2
1 parent 6444210
commit 9f472b2
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/test_core_data.py b/tests/test_core_data.py
@@ -329,8 +329,8 @@ def __init__(self, coords, velocity, metallicity, mass, age, pixel_assignment):
 
 
 class MockGasData:
-    def __init__(self, velocity=None):
-        self.velocity = velocity
+    def __init__(self, coords=None):
+        self.coords = coords
 
 
 @patch("rubix.core.data.reshape_array")
@@ -347,7 +347,7 @@ def test_get_reshape_data(mock_reshape_array):
             age=jnp.array([4.5, 5.5]),
             pixel_assignment=jnp.array([0, 1]),
         ),
-        MockGasData(velocity=None),
+        MockGasData(coords=None),
     )
 
     reshaped_data = MockRubixData(
@@ -359,7 +359,7 @@ def test_get_reshape_data(mock_reshape_array):
             age=jnp.array([4.5]),
             pixel_assignment=jnp.array([0]),
         ),
-        MockGasData(velocity=None),
+        MockGasData(coords=None),
     )
 
     def side_effect(x):