[XLA] Add test blacklist mechanism for XLA C++ unit tests.

PiperOrigin-RevId: 164124423

1 files changed, 87 insertions, 9 deletions

diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h
index 7f987a21ca..3878ac1013 100644
--- a/tensorflow/compiler/xla/tests/test_macros.h
+++ b/tensorflow/compiler/xla/tests/test_macros.h
@@ -33,13 +33,6 @@ limitations under the License.
 #include "tensorflow/compiler/xla/types.h"
 #include "tensorflow/core/platform/test.h"
 
-// Use this macro instead of directly using TEST_P for parameterized tests,
-// otherwise DISABLED_ON_* macros nested in TEST_P will not get expanded since
-// TEST_P stringifies its argument. That makes the test disabled for all targets
-// when any one of the DISABLED_ON_* macro is used, and the test will just pass.
-// TODO(b/29122096): Remove this once TEST_P fixes this problem.
-#define XLA_TEST_P(test_case_name, test_name) TEST_P(test_case_name, test_name)
-
 #define DISABLED_ON_CPU(X) X
 #define DISABLED_ON_CPU_PARALLEL(X) X
 #define DISABLED_ON_GPU(X) X
@@ -71,6 +64,91 @@ limitations under the License.
 
 // clang-format on
 
-#define XLA_TEST_F(test_fixture, test_name) TEST_F(test_fixture, test_name)
-
+namespace xla {
+
+// Reads a disabled manifest file (and retains it as a singleton) to resolve
+// whether test cases should be disabled on a particular platform.
+string PrependDisabledIfIndicated(const string& test_case_name,
+                                  const string& test_name);
+
+}  // namespace xla
+
+// This is the internal "gtest" class instantiation -- it is identical to the
+// GTEST_TEST_ macro, except that we intercept the test name for potential
+// modification by PrependDisabledIfIndicated. That file can use an arbitrary
+// heuristic to decide whether the test case should be disabled, and we
+// determine whether the test case should be disabled by resolving the (test
+// case name, test name) in a manifest file.
+#define XLA_GTEST_TEST_(test_case_name, test_name, parent_class, parent_id)   \
+  class GTEST_TEST_CLASS_NAME_(test_case_name, test_name)                     \
+      : public parent_class {                                                 \
+   public:                                                                    \
+    GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {}                    \
+                                                                              \
+   private:                                                                   \
+    virtual void TestBody();                                                  \
+    static ::testing::TestInfo* const test_info_ GTEST_ATTRIBUTE_UNUSED_;     \
+    GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name,    \
+                                                           test_name));       \
+  };                                                                          \
+                                                                              \
+  ::testing::TestInfo* const GTEST_TEST_CLASS_NAME_(test_case_name,           \
+                                                    test_name)::test_info_ =  \
+      ::testing::internal::MakeAndRegisterTestInfo(                           \
+          #test_case_name,                                                    \
+          PrependDisabledIfIndicated(#test_case_name, #test_name).c_str(),    \
+          nullptr, nullptr,                                                   \
+          ::testing::internal::CodeLocation(__FILE__, __LINE__), (parent_id), \
+          parent_class::SetUpTestCase, parent_class::TearDownTestCase,        \
+          new ::testing::internal::TestFactoryImpl<GTEST_TEST_CLASS_NAME_(    \
+              test_case_name, test_name)>);                                   \
+  void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody()
+
+// This is identical to the TEST_F macro from "gtest", but it potentially
+// disables the test based on an external manifest file, DISABLED_MANIFEST.
+//
+// Per usual, you can see what tests are available via --gunit_list_tests and
+// choose to run tests that have been disabled via the manifest via
+// --gunit_also_run_disabled_tests.
+#define XLA_TEST_F(test_fixture, test_name)              \
+  XLA_GTEST_TEST_(test_fixture, test_name, test_fixture, \
+                  ::testing::internal::GetTypeId<test_fixture>())
+
+// Likewise, this is identical to the TEST_P macro from "gtest", but
+// potentially disables the test based on the DISABLED_MANIFEST file.
+//
+// We have to wrap this in an outer layer so that any DISABLED_ON_* macros will
+// be properly expanded before the stringification occurs.
+#define XLA_TEST_P_IMPL_(test_case_name, test_name)                            \
+  class GTEST_TEST_CLASS_NAME_(test_case_name, test_name)                      \
+      : public test_case_name {                                                \
+   public:                                                                     \
+    GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {}                     \
+    virtual void TestBody();                                                   \
+                                                                               \
+   private:                                                                    \
+    static int AddToRegistry() {                                               \
+      ::testing::UnitTest::GetInstance()                                       \
+          ->parameterized_test_registry()                                      \
+          .GetTestCasePatternHolder<test_case_name>(                           \
+              #test_case_name,                                                 \
+              ::testing::internal::CodeLocation(__FILE__, __LINE__))           \
+          ->AddTestPattern(                                                    \
+              #test_case_name,                                                 \
+              PrependDisabledIfIndicated(#test_case_name, #test_name).c_str(), \
+              new ::testing::internal::TestMetaFactory<GTEST_TEST_CLASS_NAME_( \
+                  test_case_name, test_name)>());                              \
+      return 0;                                                                \
+    }                                                                          \
+    static int gtest_registering_dummy_ GTEST_ATTRIBUTE_UNUSED_;               \
+    GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name,     \
+                                                           test_name));        \
+  };                                                                           \
+  int GTEST_TEST_CLASS_NAME_(test_case_name,                                   \
+                             test_name)::gtest_registering_dummy_ =            \
+      GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::AddToRegistry();      \
+  void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody()
+
+#define XLA_TEST_P(test_case_name, test_name) \
+  XLA_TEST_P_IMPL_(test_case_name, test_name)
 #endif  // TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_