1 files changed, 24 insertions, 15 deletions
diff --git a/third_party/abseil-cpp/absl/random/internal/mock_overload_set.h b/third_party/abseil-cpp/absl/random/internal/mock_overload_set.h
index c2a30d89d5..0d9c6c120c 100644
--- a/third_party/abseil-cpp/absl/random/internal/mock_overload_set.h
+++ b/third_party/abseil-cpp/absl/random/internal/mock_overload_set.h
@@ -19,7 +19,7 @@
 #include <type_traits>
 
 #include "gmock/gmock.h"
-#include "gtest/gtest.h"
+#include "absl/random/internal/mock_helpers.h"
 #include "absl/random/mocking_bit_gen.h"
 
 namespace absl {
@@ -35,17 +35,22 @@ struct MockSingleOverload;
 // EXPECT_CALL(mock_single_overload, Call(...))` will expand to a call to
 // `mock_single_overload.gmock_Call(...)`. Because expectations are stored on
 // the MockingBitGen (an argument passed inside `Call(...)`), this forwards to
-// arguments to Mocking::Register.
+// arguments to MockingBitGen::Register.
+//
+// The underlying KeyT must match the KeyT constructed by DistributionCaller.
 template <typename DistrT, typename Ret, typename... Args>
 struct MockSingleOverload<DistrT, Ret(MockingBitGen&, Args...)> {
   static_assert(std::is_same<typename DistrT::result_type, Ret>::value,
                 "Overload signature must have return type matching the "
-                "distributions result type.");
-  auto gmock_Call(
-      absl::MockingBitGen& gen,  // NOLINT(google-runtime-references)
-      const ::testing::Matcher<Args>&... args)
-      -> decltype(gen.Register<DistrT, Args...>(args...)) {
-    return gen.Register<DistrT, Args...>(args...);
+                "distribution result_type.");
+  using KeyT = Ret(DistrT, std::tuple<Args...>);
+
+  template <typename MockURBG>
+  auto gmock_Call(MockURBG& gen, const ::testing::Matcher<Args>&... matchers)
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...)) {
+    static_assert(std::is_base_of<MockingBitGen, MockURBG>::value,
+                  "Mocking requires an absl::MockingBitGen");
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...);
   }
 };
 
@@ -53,13 +58,17 @@ template <typename DistrT, typename Ret, typename Arg, typename... Args>
 struct MockSingleOverload<DistrT, Ret(Arg, MockingBitGen&, Args...)> {
   static_assert(std::is_same<typename DistrT::result_type, Ret>::value,
                 "Overload signature must have return type matching the "
-                "distributions result type.");
-  auto gmock_Call(
-      const ::testing::Matcher<Arg>& arg,
-      absl::MockingBitGen& gen,  // NOLINT(google-runtime-references)
-      const ::testing::Matcher<Args>&... args)
-      -> decltype(gen.Register<DistrT, Arg, Args...>(arg, args...)) {
-    return gen.Register<DistrT, Arg, Args...>(arg, args...);
+                "distribution result_type.");
+  using KeyT = Ret(DistrT, std::tuple<Arg, Args...>);
+
+  template <typename MockURBG>
+  auto gmock_Call(const ::testing::Matcher<Arg>& matcher, MockURBG& gen,
+                  const ::testing::Matcher<Args>&... matchers)
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher,
+                                                             matchers...)) {
+    static_assert(std::is_base_of<MockingBitGen, MockURBG>::value,
+                  "Mocking requires an absl::MockingBitGen");
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher, matchers...);
   }
 };