From 5cd30374430fe1991e6f2b22a07f817a88b8a0bf Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Mon, 20 Nov 2023 03:00:09 -0800
Subject: [PATCH] Adding a GpuTestWithParamBase test class to support value
 parameterized tests

PiperOrigin-RevId: 583967017
---
 mediapipe/gpu/gpu_test_base.h | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/mediapipe/gpu/gpu_test_base.h b/mediapipe/gpu/gpu_test_base.h
index 6ec53603b..94842c769 100644
--- a/mediapipe/gpu/gpu_test_base.h
+++ b/mediapipe/gpu/gpu_test_base.h
@@ -15,6 +15,9 @@
 #ifndef MEDIAPIPE_GPU_GPU_TEST_BASE_H_
 #define MEDIAPIPE_GPU_GPU_TEST_BASE_H_
 
+#include <functional>
+#include <memory>
+
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
 #include "mediapipe/gpu/gl_calculator_helper.h"
@@ -22,9 +25,9 @@
 
 namespace mediapipe {
 
-class GpuTestBase : public ::testing::Test {
+class GpuTestEnvironment {
  protected:
-  GpuTestBase() { helper_.InitializeForTest(gpu_resources_.get()); }
+  GpuTestEnvironment() { helper_.InitializeForTest(gpu_resources_.get()); }
 
   void RunInGlContext(std::function<void(void)> gl_func) {
     helper_.RunInGlContext(std::move(gl_func));
@@ -35,6 +38,12 @@ class GpuTestBase : public ::testing::Test {
   GlCalculatorHelper helper_;
 };
 
+class GpuTestBase : public testing::Test, public GpuTestEnvironment {};
+
+template <typename T>
+class GpuTestWithParamBase : public testing::TestWithParam<T>,
+                             public GpuTestEnvironment {};
+
 }  // namespace mediapipe
 
 #endif  // MEDIAPIPE_GPU_GPU_TEST_BASE_H_