From b11178ca5da8181bdd3af2afb85ac0d11db9eaf4 Mon Sep 17 00:00:00 2001
From: "Haoyu (Daniel)" <yanghaoyu97@outlook.com>
Date: Sat, 28 Sep 2024 11:37:45 +0800
Subject: [PATCH] add unit test for lazy mp api key

---
 tests/test_mprester.py | 36 ++++++++++++++++++++++++++----------
 1 file changed, 26 insertions(+), 10 deletions(-)

diff --git a/tests/test_mprester.py b/tests/test_mprester.py
index 4b7edb82..bd8a4e81 100644
--- a/tests/test_mprester.py
+++ b/tests/test_mprester.py
@@ -1,7 +1,7 @@
 import itertools
 import os
 import random
-from unittest.mock import patch
+import importlib
 
 import numpy as np
 import pytest
@@ -28,13 +28,15 @@
 from mp_api.client import MPRester
 from mp_api.client.core.settings import MAPIClientSettings
 
+FAKE_MP_API_KEY = "12345678901234567890123456789012"  # 32 chars
+
 
 @pytest.fixture()
 def mpr():
     # Only mock the API key in GitHub CI environment,
     # otherwise people cannot run tests locally as the key is invalid
     if os.getenv("GITHUB_ACTIONS") is not None:
-        os.environ["MP_API_KEY"] = "12345678901234567890123456789012"
+        os.environ["MP_API_KEY"] = FAKE_MP_API_KEY
 
     rester = MPRester()
     yield rester
@@ -331,14 +333,28 @@ def test_large_list(self, mpr):
         docs = mpr.summary.search(material_ids=mpids, fields=["material_ids"])
         assert len(docs) == 10000
 
+    def test_get_api_key_from_env_var(self, monkeypatch: pytest.MonkeyPatch):
+        """Ensure the MP_API_KEY from environment variable is retrieved
+        at runtime, not import time.
+        """
+        # Mock a invalid key set before import MPRester
+        import mp_api.client
+
+        monkeypatch.setenv("MP_API_KEY", "INVALID KEY")
+
+        importlib.reload(mp_api.client)
+
+        monkeypatch.setenv("MP_API_KEY", FAKE_MP_API_KEY)
+        assert MPRester().api_key == FAKE_MP_API_KEY
 
-@patch.dict(os.environ, {"MP_API_KEY": ""})  # Unset MP_API_KEY
-def test_get_api_key_from_settings(monkeypatch: pytest.MonkeyPatch):
-    """Test getting MP_API_KEY from SETTINGS."""
-    fake_api_key = "12345678901234567890123456789012"  # 32 chars
+    def test_get_api_key_from_settings(self, monkeypatch: pytest.MonkeyPatch):
+        """Test environment variable "MP_API_KEY" is not set and
+        get "PMG_MAPI_KEY" from "SETTINGS".
+        """
+        monkeypatch.delenv("MP_API_KEY", raising=False)
 
-    # patch pymatgen.core.SETTINGS to contain PMG_MAPI_KEY
-    monkeypatch.setitem(SETTINGS, "PMG_MAPI_KEY", fake_api_key)
+        # patch pymatgen.core.SETTINGS to contain PMG_MAPI_KEY
+        monkeypatch.setitem(SETTINGS, "PMG_MAPI_KEY", FAKE_MP_API_KEY)
 
-    # create MPRester and check that it picked up the API key from pymatgen SETTINGS
-    assert MPRester().api_key == fake_api_key
+        # create MPRester and check that it picked up the API key from SETTINGS
+        assert MPRester().api_key == FAKE_MP_API_KEY