From b0ae3a8374ae54eec846e1f422da554e792acc01 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Mon, 6 Apr 2026 07:32:23 +0000
Subject: [PATCH] Add mock unit tests for Materials Project integration

This commit adds `tests/test_materialsproject.py`, which provides
unit tests for the `search` and `by_id` functions in
`src/structuretoolkit/build/materialsproject.py`.

The tests use `unittest.mock` to simulate the `mp_api.client.MPRester`
and `structuretoolkit.build.materialsproject.pymatgen_to_ase`
functions, ensuring that the integration logic is correctly tested
without requiring a real API key or an internet connection.

Test cases include:
- `test_search`: Verifies that `materialsproject_search` yields processed
  results and correctly handles the `structure` field.
- `test_by_id`: Verifies that `materialsproject_by_id` correctly returns
  a single structure for `final=True` and a list of structures for
  `final=False`.

The `mp_api` module is mocked in `sys.modules` to avoid import errors
in environments where it is not installed.

Co-authored-by: jan-janssen <3854739+jan-janssen@users.noreply.github.com>
---
 tests/test_materialsproject.py | 65 ++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)
 create mode 100644 tests/test_materialsproject.py

diff --git a/tests/test_materialsproject.py b/tests/test_materialsproject.py
new file mode 100644
index 000000000..dd15f6b17
--- /dev/null
+++ b/tests/test_materialsproject.py
@@ -0,0 +1,65 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import sys
+
+# Mock mp_api before importing structuretoolkit
+mock_mp_api = MagicMock()
+sys.modules["mp_api"] = mock_mp_api
+sys.modules["mp_api.client"] = mock_mp_api.client
+
+import structuretoolkit as stk
+
+class TestMaterialsProject(unittest.TestCase):
+    @patch("mp_api.client.MPRester")
+    @patch("structuretoolkit.build.materialsproject.pymatgen_to_ase")
+    def test_search(self, mock_pymatgen_to_ase, mock_mp_rester):
+        # Setup mock for MPRester as a context manager
+        mock_mpr = MagicMock()
+        mock_mp_rester.return_value.__enter__.return_value = mock_mpr
+
+        # Setup mock for summary.search
+        mock_mpr.summary.search.return_value = [
+            {"material_id": "mp-1", "structure": "mock_pmg_struct"}
+        ]
+
+        # Setup mock for pymatgen_to_ase
+        mock_pymatgen_to_ase.return_value = "mock_ase_struct"
+
+        # Call search
+        results = list(stk.build.materialsproject_search("Fe"))
+
+        # Assertions
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]["material_id"], "mp-1")
+        self.assertEqual(results[0]["structure"], "mock_ase_struct")
+        mock_mpr.summary.search.assert_called_once()
+        mock_pymatgen_to_ase.assert_called_once_with("mock_pmg_struct")
+
+    @patch("mp_api.client.MPRester")
+    @patch("structuretoolkit.build.materialsproject.pymatgen_to_ase")
+    def test_by_id(self, mock_pymatgen_to_ase, mock_mp_rester):
+        # Setup mock for MPRester as a context manager
+        mock_mpr = MagicMock()
+        mock_mp_rester.return_value.__enter__.return_value = mock_mpr
+
+        # Setup mock for pymatgen_to_ase
+        mock_pymatgen_to_ase.side_effect = lambda x: f"ase_{x}"
+
+        # Test final=True
+        mock_mpr.get_structure_by_material_id.return_value = "pmg_struct"
+        res = stk.build.materialsproject_by_id("mp-1", final=True)
+        self.assertEqual(res, "ase_pmg_struct")
+        mock_mpr.get_structure_by_material_id.assert_called_with(
+            material_id="mp-1", final=True, conventional_unit_cell=False
+        )
+
+        # Test final=False
+        mock_mpr.get_structure_by_material_id.return_value = ["pmg_1", "pmg_2"]
+        res = stk.build.materialsproject_by_id("mp-1", final=False)
+        self.assertEqual(res, ["ase_pmg_1", "ase_pmg_2"])
+        mock_mpr.get_structure_by_material_id.assert_called_with(
+            material_id="mp-1", final=False, conventional_unit_cell=False
+        )
+
+if __name__ == "__main__":
+    unittest.main()