1 files changed, 16 insertions, 10 deletions
diff --git a/absl/random/internal/mock_overload_set.h b/absl/random/internal/mock_overload_set.h
index c2a30d89..dccc6cee 100644
--- a/absl/random/internal/mock_overload_set.h
+++ b/absl/random/internal/mock_overload_set.h
@@ -20,6 +20,7 @@
 
 #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 +36,20 @@ 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.");
+                "distribution result_type.");
+  using KeyT = Ret(DistrT, std::tuple<Args...>);
   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...);
+      const ::testing::Matcher<Args>&... matchers)
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...)) {
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...);
   }
 };
 
@@ -53,13 +57,15 @@ 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.");
+                "distribution result_type.");
+  using KeyT = Ret(DistrT, std::tuple<Arg, Args...>);
   auto gmock_Call(
-      const ::testing::Matcher<Arg>& arg,
+      const ::testing::Matcher<Arg>& matcher,
       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...);
+      const ::testing::Matcher<Args>&... matchers)
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher,
+                                                             matchers...)) {
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher, matchers...);
   }
 };