test fixes

Author: esoteric-ephemera

mp_api/_test_utils.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+from enum import Enum
+
 try:
     import pytest
 except ImportError as exc:
@@ -108,19 +110,40 @@ def client_pagination(
     ) == set()
 
 
-def client_sort(search_method: Callable, sort_fields: str | Sequence[str]):
+def client_sort(
+    search_method: Callable,
+    sort_fields: str | Sequence[str],
+    aux_query: dict[str, Any] | None = None,
+):
     """Test sorting on an endpoint.
 
     Args:
     search_method (Callable) : Client search method to use
     sort_fields (str or Sequence of str) : fields to sort on
+    aux_query (dict) : auxiliary query needed to filter documents
 
     Raises:
     AssertionError if sorting in ascending or descending order does not work.
     """
+
+    def _normalize(doc, field: str):
+        v = getattr(doc, field)
+        # serialize enums
+        return v.value if isinstance(v, Enum) else v
+
+    user_query = {
+        k: v
+        for k, v in (aux_query or {}).items()
+        if k not in ("_page", "_sort_fields", "chunk_size", "fields")
+    }
     for sort_field in [sort_fields] if isinstance(sort_fields, str) else sort_fields:
+
         asc = search_method(
-            _page=1, _sort_fields=sort_field, chunk_size=NUM_DOCS, fields=[sort_field]
+            _page=1,
+            _sort_fields=sort_field,
+            chunk_size=NUM_DOCS,
+            fields=[sort_field, "deprecated", "material_id"],
+            **user_query,
         )
         desc = search_method(
             _page=1,
@@ -130,12 +153,12 @@ def client_sort(search_method: Callable, sort_fields: str | Sequence[str]):
         )
 
         idxs = list(range(NUM_DOCS))
-        assert sorted(idxs, key=lambda idx: getattr(asc[idx], sort_field)) == idxs
+        assert sorted(idxs, key=lambda idx: _normalize(asc[idx], sort_field)) == idxs
 
         assert (
             sorted(
                 idxs,
-                key=lambda idx: getattr(desc[idx], sort_field),
+                key=lambda idx: _normalize(desc[idx], sort_field),
                 reverse=True,
             )
             == idxs

mp_api/client/core/utils.py

@@ -13,7 +13,7 @@
 import pyarrow.dataset as ds
 from deltalake import DeltaTable
 from emmet.core import __version__ as _EMMET_CORE_VER
-from emmet.core.mpid_ext import validate_identifier
+from emmet.core.mpid import validate_identifier
 from monty.json import MontyDecoder
 from packaging.version import parse as parse_version
 

mp_api/client/routes/materials/phonon.py

@@ -225,7 +225,7 @@ def compute_thermo_quantities(
 
         self.use_document_model = True
         docs[0]["phonon_dos"] = self.get_dos_from_phonon_id(  # type: ignore[index]
-            phonon_id, phonon_method
+            phonon_id, phonon_method  # type: ignore[arg-type]
         )
         doc = PhononBSDOSDoc(**docs[0])  # type: ignore[arg-type]
         self.use_document_model = use_document_model

mp_api/client/routes/materials/thermo.py

@@ -18,7 +18,7 @@
 class ThermoRester(BaseRester):
     suffix = "materials/thermo"
     document_model = ThermoDoc  # type: ignore
-    primary_key = "thermo_id"
+    primary_key = "material_id"
 
     def search(
         self,

mp_api/client/routes/materials/xas.py

@@ -17,7 +17,7 @@
 class XASRester(BaseRester):
     suffix = "materials/xas"
     document_model = XASDoc  # type: ignore
-    primary_key = "spectrum_id"
+    primary_key = "material_id"
     delta_backed = False
 
     def search(

tests/client/materials/test_electrodes.py

@@ -95,11 +95,8 @@ def test_pagination():
         client_pagination(rester.search, "material_ids")
 
 
-@pytest.mark.xfail(reason="Sort requires API redeployment", strict=False)
 @requires_api_key
-@pytest.mark.parametrize(
-    "sort_field", ["working_ion", "stability_charge", "average_voltage"]
-)
+@pytest.mark.parametrize("sort_field", ["stability_charge", "average_voltage"])
 def test_sort(sort_field):
     with ElectrodeRester() as rester:
         client_sort(rester.search, sort_field)

tests/client/materials/test_provenance.py

@@ -25,7 +25,7 @@ def rester():
 
 alt_name_dict: dict = {"material_ids": "material_id"}
 
-custom_field_tests: dict = {"material_ids": ["mp-149"]}
+custom_field_tests: dict = {"material_ids": ["mp-13"]}
 
 
 @requires_api_key

tests/client/materials/test_summary.py

@@ -161,4 +161,4 @@ def test_pagination():
 @pytest.mark.parametrize("sort_field", summary_sort_fields)
 def test_sort(sort_field: str):
     with SummaryRester() as rester:
-        client_sort(rester.search, sort_field)
+        client_sort(rester.search, sort_field, aux_query={sort_field: (0, 10)})

tests/client/materials/test_xas.py

@@ -64,7 +64,6 @@ def test_client(rester):
     )
 
 
-# TODO: how to test pagination now that spectrum_id is computed, not stored?
 @requires_api_key
 def test_pagination():
     with XASRester() as rester:
@@ -73,10 +72,13 @@ def test_pagination():
         )
 
 
-@pytest.mark.xfail(reason="Sort requires API redeployment", strict=False)
 @requires_api_key
 @pytest.mark.parametrize(
-    "sort_field", ["spectrum_type", "absorbing_element", "chemsys"]
+    "sort_field",
+    [
+        "spectrum_type",
+        "absorbing_element",
+    ],
 )
 def test_sort(sort_field):
     with XASRester() as rester:

tests/client/molecules/test_jcesr.py

@@ -68,9 +68,13 @@ def test_pagination():
         client_pagination(rester.search, "task_id")
 
 
-@pytest.mark.xfail(reason="Sort requires API redeployment", strict=False)
 @requires_api_key
-@pytest.mark.parametrize("sort_field", ["task_id", "IE", "EA"])
+@pytest.mark.parametrize(
+    "sort_field",
+    [
+        "task_id",
+    ],
+)
 def test_sort(sort_field):
     with JcesrMoleculesRester() as rester:
         client_sort(rester.search, sort_field)