From b6dc30dd14aa4e540dceebcab7891db25e035dcd Mon Sep 17 00:00:00 2001
From: Hai Nguyen <hainm.comp@gmail.com>
Date: Fri, 22 Nov 2024 13:45:39 -0500
Subject: [PATCH] more

---
 nglview/tests/test_widget.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/nglview/tests/test_widget.py b/nglview/tests/test_widget.py
index bdf5ee1b..6e056ecb 100644
--- 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():