1 files changed, 23 insertions, 10 deletions
diff --git a/absl/random/internal/mock_overload_set.h b/absl/random/internal/mock_overload_set.h
index c2a30d89..60561b51 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/base/internal/fast_type_id.h"
 #include "absl/random/mocking_bit_gen.h"
 
 namespace absl {
@@ -35,17 +36,23 @@ 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.
 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 ArgTupleT = 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(gen.RegisterMock<Ret, ArgTupleT>(
+                         std::declval<::absl::base_internal::FastTypeIdType>())
+                      .gmock_Call(matchers...)) {
+    return gen
+        .RegisterMock<Ret, ArgTupleT>(
+            ::absl::base_internal::FastTypeId<Ret(DistrT, ArgTupleT)>())
+        .gmock_Call(matchers...);
   }
 };
 
@@ -53,13 +60,19 @@ 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 ArgTupleT = 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(gen.RegisterMock<Ret, ArgTupleT>(
+                         std::declval<::absl::base_internal::FastTypeIdType>())
+                      .gmock_Call(matcher, matchers...)) {
+    return gen
+        .RegisterMock<Ret, ArgTupleT>(
+            ::absl::base_internal::FastTypeId<Ret(DistrT, ArgTupleT)>())
+        .gmock_Call(matcher, matchers...);
   }
 };