more

nglviewer · Nov 22, 2024 · b6dc30d · b6dc30d
1 parent 9fd4975
commit b6dc30d
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/nglview/tests/test_widget.py b/nglview/tests/test_widget.py
@@ -88,10 +88,14 @@
 
 
 def get_mocked_traj():
-    traj = MagicMock()
-    traj.get_coordinates.return_value = np.random.rand(10, 3, 3)
-    traj.n_frames = 10
-    return traj
+    class MockedTraj(nv.Structure):
+        def __init__(self):
+            self.n_frames = 10
+
+        def get_coordinates(self, frame):
+            return np.random.rand(10, 3, 3)
+
+    return MockedTraj()
 
 
 def default_view():