ensure find_structure works with doc model false

Author: esoteric-ephemera

mp_api/client/routes/materials/materials.py

@@ -242,10 +242,13 @@ def find_structure(
             comparator=ElementComparator(),
         )
 
-        matches = []
+        matches: list[str] = []
         for doc in mat_docs:
-            if matcher.fit(s, doc.structure):
-                matches.append(doc.material_id)
+            if matcher.fit(
+                s,
+                doc.structure if self.use_document_model else Structure.from_dict(doc["structure"]),  # type: ignore[union-attr]
+            ):
+                matches.append(doc.material_id.string if self.use_document_model else doc["material_id"])  # type: ignore[union-attr]
                 if len(matches) >= max_matches:
                     break
 

tests/client/test_mprester.py

@@ -102,7 +102,9 @@ def test_get_structures(self, mpr):
         structs = mpr.get_structures("Mn-O", final=False)
         assert len(structs) > 0
 
-    def test_find_structure(self, mpr):
+    def test_find_structure(
+        self,
+    ):
         cif_str = """# mp-111
 data_Ne
 _symmetry_space_group_name_H-M   'P 1'
@@ -139,11 +141,12 @@ def test_find_structure(self, mpr):
             f.write(cif_str)
             f.seek(0)
 
-        for struct_or_path in (
-            temp_file.name,
-            struct_from_cif,
-        ):
-            data = mpr.find_structure(struct_or_path)
+        for struct_or_path, use_document_model in [
+            (temp_file.name, True),
+            (struct_from_cif, False),
+        ]:
+            with MPRester(use_document_model=use_document_model) as mpr:
+                data = mpr.find_structure(struct_or_path)
             assert isinstance(data, str) and data == "mp-111"
 
         f.close()