fix get_trajectory helper func + test

Author: tsmathis

mp_api/client/routes/materials/tasks.py

@@ -24,7 +24,7 @@ class TaskRester(BaseRester):
     primary_key: str = "task_id"
     delta_backed = True
 
-    def get_trajectory(self, task_id: MPID | AlphaID | str) -> list[dict[str, Any]]:
+    def get_trajectory(self, task_id: MPID | AlphaID | str) -> dict[str, Any]:
         """Returns a Trajectory object containing the geometry of the
         material throughout a calculation. This is most useful for
         observing how a material relaxes during a geometry optimization.
@@ -33,7 +33,7 @@ def get_trajectory(self, task_id: MPID | AlphaID | str) -> list[dict[str, Any]]:
             task_id (str, MPID, AlphaID): Task ID
 
         Returns:
-            list of dict representing emmet.core.trajectory.Trajectory
+            dict representing emmet.core.trajectory.RelaxTrajectory
         """
         as_alpha = str(AlphaID(task_id, padlen=8)).split("-")[-1]
         traj_tbl = DeltaTable(
@@ -57,7 +57,7 @@ def get_trajectory(self, task_id: MPID | AlphaID | str) -> list[dict[str, Any]]:
         if not traj_data:
             raise MPRestError(f"No trajectory data for {task_id} found")
 
-        return RelaxTrajectory(**traj_data[0]).to_pmg().as_dict()
+        return RelaxTrajectory(**traj_data[0]).model_dump()
 
     def search(
         self,

tests/client/materials/test_tasks.py

@@ -1,13 +1,14 @@
 import os
-from ..conftest import client_search_testing, requires_api_key
-import pytest
 
+import pytest
 from emmet.core.mpid import MPID, AlphaID
-from emmet.core.trajectory import Trajectory
+from emmet.core.trajectory import RelaxTrajectory
 from emmet.core.utils import utcnow
 
 from mp_api.client.routes.materials.tasks import TaskRester
 
+from ..conftest import client_search_testing, requires_api_key
+
 
 @pytest.fixture
 def rester():
@@ -57,11 +58,11 @@ def test_client(rester):
 
 @pytest.mark.parametrize("mpid", ["mp-149", MPID("mp-149"), AlphaID("mp-149")])
 def test_get_trajectories(rester, mpid):
-    trajectories = [traj for traj in rester.get_trajectory(mpid)]
+    trajectory = rester.get_trajectory(mpid)
 
     expected_model_fields = {
         field_name
-        for field_name, field in Trajectory.model_fields.items()
+        for field_name, field in RelaxTrajectory.model_fields.items()
         if not field.exclude
     }
-    assert all(set(traj) == expected_model_fields for traj in trajectories)
+    assert set(trajectory) == expected_model_fields