Fix traj test

materialsproject · May 10, 2023 · 307e9db · 307e9db
1 parent 6e225dc
commit 307e9db
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
@@ -2,14 +2,13 @@
 import typing
 
 import pytest
-from pymatgen.core.trajectory import Trajectory
 
 from mp_api.client.routes.tasks import TaskRester
 
 
 @pytest.fixture
 def rester():
-    rester = TaskRester()
+    rester = TaskRester(monty_decode=False)
     yield rester
     rester.session.close()
 
@@ -97,7 +96,7 @@ def test_client(rester):
                         ],
                     }
 
-                doc = search_method(**q)[0].dict()
+                doc = search_method(**q)[0]
                 for sub_field in sub_doc_fields:
                     if sub_field in doc:
                         doc = doc[sub_field]
@@ -109,7 +108,7 @@ def test_client(rester):
 
 
 def test_get_trajectories(rester):
-    trajectories = rester.get_trajectory("mp-149")
+    trajectories = [traj for traj in rester.get_trajectory("mp-149")]
 
     for traj in trajectories:
-        assert isinstance(traj, Trajectory)
+        assert ("@module", "pymatgen.core.trajectory") in traj.items()