Generalize constraints on MPI rank counts for tests

[alexxy/gromacs.git] / src / testutils / tests / mpitest.cpp

diff --git a/src/testutils/tests/mpitest.cpp b/src/testutils/tests/mpitest.cpp
index d41ee8d9281162648e952c26a2a47cea4f06d757..f3d8c0f4eca09bef4bea9c4eaa7090aec062614b 100644 (file)
--- a/src/testutils/tests/mpitest.cpp
+++ b/src/testutils/tests/mpitest.cpp
@@ -1,7 +1,7 @@
 /*
  * This file is part of the GROMACS molecular simulation package.
  *
- * Copyright (c) 2016,2019, by the GROMACS development team, led by
+ * Copyright (c) 2016,2019,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -45,37 +45,43 @@
 
 #include "config.h"
 
+#include <vector>
+
 #include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
 #include "gromacs/utility/basenetwork.h"
 #include "gromacs/utility/gmxmpi.h"
 
+namespace gmx
+{
+namespace test
+{
 namespace
 {
 
 class MpiSelfTest : public ::testing::Test
 {
 public:
-    MpiSelfTest() : reached{ 0, 0 } {}
-
-    int reached[2];
+    //! Whether each rank participated, relevant only on rank 0
+    std::vector<int> reached_;
 };
 
 TEST_F(MpiSelfTest, Runs)
 {
-    GMX_MPI_TEST(2);
-#if GMX_THREAD_MPI
-    reached[gmx_node_rank()] = 1;
-    MPI_Barrier(MPI_COMM_WORLD);
-#else
+    GMX_MPI_TEST(RequireMinimumRankCount<2>);
+    if (gmx_node_rank() == 0)
+    {
+        reached_.resize(getNumberOfTestMpiRanks(), 0);
+    }
     int value = 1;
-    MPI_Gather(&value, 1, MPI_INT, reached, 1, MPI_INT, 0, MPI_COMM_WORLD);
-#endif
+    MPI_Gather(&value, 1, MPI_INT, reached_.data(), 1, MPI_INT, 0, MPI_COMM_WORLD);
     if (gmx_node_rank() == 0)
     {
-        EXPECT_EQ(1, reached[0]);
-        EXPECT_EQ(1, reached[1]);
+        EXPECT_THAT(reached_, testing::Each(value));
     }
 }
 
 } // namespace
+} // namespace test
+} // namespace gmx