fix bug

Author: Han Wang

source/tests/pt/test_env_threads.py

@@ -7,7 +7,7 @@
 import deepmd.env as common_env
 
 
-def test_env_threads_guard_handles_runtimeerror(monkeypatch, caplog) -> None:
+def test_env_threads_guard_handles_runtimeerror(monkeypatch) -> None:
     def raise_err(*_args, **_kwargs) -> None:
         raise RuntimeError("boom")
 
@@ -18,11 +18,17 @@ def raise_err(*_args, **_kwargs) -> None:
     monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
     monkeypatch.setattr(torch, "set_num_threads", raise_err)
 
-    caplog.set_level(logging.WARNING, logger="deepmd.pt.utils.env")
+    messages: list[str] = []
+    original_warning = logging.Logger.warning
+
+    def capture_warning(self, msg, *args, **kwargs):  # type: ignore[no-untyped-def]
+        messages.append(str(msg))
+        return original_warning(self, msg, *args, **kwargs)
+
+    monkeypatch.setattr(logging.Logger, "warning", capture_warning)
     import deepmd.pt.utils.env as env
 
     importlib.reload(env)
 
-    messages = [record.getMessage() for record in caplog.records]
     assert any("Could not set torch interop threads" in msg for msg in messages)
     assert any("Could not set torch intra threads" in msg for msg in messages)

source/tests/pt_expt/utils/test_env.py

@@ -7,7 +7,7 @@
 import deepmd.env as common_env
 
 
-def test_env_threads_guard_handles_runtimeerror(monkeypatch, caplog) -> None:
+def test_env_threads_guard_handles_runtimeerror(monkeypatch) -> None:
     def raise_err(*_args, **_kwargs) -> None:
         raise RuntimeError("boom")
 
@@ -18,11 +18,17 @@ def raise_err(*_args, **_kwargs) -> None:
     monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
     monkeypatch.setattr(torch, "set_num_threads", raise_err)
 
-    caplog.set_level(logging.WARNING, logger="deepmd.pt_expt.utils.env")
+    messages: list[str] = []
+    original_warning = logging.Logger.warning
+
+    def capture_warning(self, msg, *args, **kwargs):  # type: ignore[no-untyped-def]
+        messages.append(str(msg))
+        return original_warning(self, msg, *args, **kwargs)
+
+    monkeypatch.setattr(logging.Logger, "warning", capture_warning)
     import deepmd.pt_expt.utils.env as env
 
     importlib.reload(env)
 
-    messages = [record.getMessage() for record in caplog.records]
     assert any("Could not set torch interop threads" in msg for msg in messages)
     assert any("Could not set torch intra threads" in msg for msg in messages)