From 8890b315cf837a65bc8a07bf764b85603df492c1 Mon Sep 17 00:00:00 2001
From: qbisi <qbisicwate@gmail.com>
Date: Tue, 7 Oct 2025 16:40:29 +0800
Subject: [PATCH 1/2] python3Packages.mpi4py: add test for mpich

---
 pkgs/development/python-modules/mpi4py/default.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pkgs/development/python-modules/mpi4py/default.nix b/pkgs/development/python-modules/mpi4py/default.nix
index 42422d61a0d7..754d747394eb 100644
--- a/pkgs/development/python-modules/mpi4py/default.nix
+++ b/pkgs/development/python-modules/mpi4py/default.nix
@@ -8,6 +8,8 @@
   toPythonModule,
   pytestCheckHook,
   mpiCheckPhaseHook,
+  mpi4py,
+  mpich,
 }:
 
 buildPythonPackage rec {
@@ -60,6 +62,10 @@ buildPythonPackage rec {
 
   passthru = {
     inherit mpi;
+
+    tests = {
+      mpich = mpi4py.override { mpi = mpich; };
+    };
   };
 
   meta = {

From 47fffffa17d6de1139aa114bf61d1d3415c29bbe Mon Sep 17 00:00:00 2001
From: qbisi <qbisicwate@gmail.com>
Date: Tue, 7 Oct 2025 17:21:55 +0800
Subject: [PATCH 2/2] python3Packages.mpi4py: follow upstream checkPhase

---
 .../python-modules/mpi4py/default.nix         | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/pkgs/development/python-modules/mpi4py/default.nix b/pkgs/development/python-modules/mpi4py/default.nix
index 754d747394eb..93162a1bff7f 100644
--- a/pkgs/development/python-modules/mpi4py/default.nix
+++ b/pkgs/development/python-modules/mpi4py/default.nix
@@ -6,7 +6,7 @@
   setuptools,
   mpi,
   toPythonModule,
-  pytestCheckHook,
+  pytest,
   mpiCheckPhaseHook,
   mpi4py,
   mpich,
@@ -42,17 +42,9 @@ buildPythonPackage rec {
   pythonImportsCheck = [ "mpi4py" ];
 
   nativeCheckInputs = [
-    pytestCheckHook
+    pytest
     mpiCheckPhaseHook
   ];
-  disabledTestPaths = lib.optionals (mpi.pname == "mpich") [
-    # These tests from some reason cause pytest to crash, and therefor it is
-    # hard to debug them. Upstream mentions these tests to raise issues in
-    # https://github.com/mpi4py/mpi4py/issues/418  but the workaround suggested
-    # there (setting MPI4PY_RC_RECV_MPROBE=0) doesn't work.
-    "test/test_util_pool.py"
-    "demo/futures/test_futures.py"
-  ];
 
   __darwinAllowLocalNetworking = true;
 
@@ -60,6 +52,19 @@ buildPythonPackage rec {
   # see https://github.com/mpi4py/mpi4py/issues/545#issuecomment-2343011460
   env.MPI4PY_TEST_SPAWN = if mpi.pname == "openmpi" then 0 else 1;
 
+  # follow upstream's checkPhase
+  # see https://github.com/mpi4py/mpi4py/blob/4.1.0/.github/workflows/ci-test.yml#L92-L95
+  checkPhase = ''
+    runHook preCheck
+
+    echo 'Testing mpi4py (np=1)'
+    mpiexec -n 1 python test/main.py -v
+    echo 'Testing mpi4py (np=2)'
+    mpiexec -n 2 python test/main.py -v -f -e spawn
+
+    runHook postCheck
+  '';
+
   passthru = {
     inherit mpi;