Fix trajectory convenience function (#542)

* Fix trajectory convenience function

* Add trajectory test

* Fix task test

* Test bug fix

src/mp_api/routes/tasks.py

@@ -1,7 +1,7 @@
 from typing import List, Optional, Union
 
 from emmet.core.tasks import TaskDoc
-from mp_api.core.client import BaseRester
+from mp_api.core.client import BaseRester, MPRestError
 
 
 class TaskRester(BaseRester[TaskDoc]):
@@ -17,10 +17,17 @@ def get_trajectory(self, task_id):
         observing how a material relaxes during a geometry optimization.
 
         :param task_id: A specified task_id
-        :return: Trajectory object
+        :return: List of trajectory objects
         """
 
-        pass
+        traj_data = self._query_resource_data(
+            suburl=f"trajectory/{task_id}/", use_document_model=False
+        )[0].get("trajectories", None)
+
+        if traj_data is None:
+            raise MPRestError(f"No trajectory data for {task_id} found")
+
+        return traj_data
 
     def search_task_docs(
         self,
@@ -65,5 +72,5 @@ def search_task_docs(
             chunk_size=chunk_size,
             all_fields=all_fields,
             fields=fields,
-            **query_params
+            **query_params,
         )

tests/test_tasks.py

@@ -1,5 +1,6 @@
 import os
 import pytest
+from pymatgen.core.trajectory import Trajectory
 from mp_api.routes.tasks import TaskRester
 
 import inspect
@@ -94,3 +95,11 @@ def test_client(rester):
                     doc[project_field if project_field is not None else param]
                     is not None
                 )
+
+
+def test_get_trajectories():
+
+    trajectories = resters[0].get_trajectory("mp-149")
+
+    for traj in trajectories:
+        assert isinstance(traj, Trajectory)