diff --git a/xla/pjrt/plugin/xla_tpu/BUILD b/xla/pjrt/plugin/xla_tpu/BUILD
index e6f08e70b98650..85a27a30e2511b 100644
--- a/xla/pjrt/plugin/xla_tpu/BUILD
+++ b/xla/pjrt/plugin/xla_tpu/BUILD
@@ -33,6 +33,7 @@ cc_test(
         "//xla/pjrt:pjrt_common",
         "//xla/tests:xla_internal_test_main",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:test",
     ],
 )
diff --git a/xla/pjrt/plugin/xla_tpu/xla_tpu_pjrt_client_test.cc b/xla/pjrt/plugin/xla_tpu/xla_tpu_pjrt_client_test.cc
index 7c7b8c587d45bb..4148cca11e650f 100644
--- a/xla/pjrt/plugin/xla_tpu/xla_tpu_pjrt_client_test.cc
+++ b/xla/pjrt/plugin/xla_tpu/xla_tpu_pjrt_client_test.cc
@@ -17,6 +17,8 @@ limitations under the License.
 
 #include <string>
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
 #include "absl/container/flat_hash_map.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "tsl/platform/test.h"