1 files changed, 10 insertions, 17 deletions

diff --git a/absl/random/internal/mock_overload_set.h b/absl/random/internal/mock_overload_set.h
index 60561b51..dccc6cee 100644
--- a/absl/random/internal/mock_overload_set.h
+++ b/absl/random/internal/mock_overload_set.h
@@ -20,7 +20,7 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
-#include "absl/base/internal/fast_type_id.h"
+#include "absl/random/internal/mock_helpers.h"
 #include "absl/random/mocking_bit_gen.h"
 
 namespace absl {
@@ -37,22 +37,19 @@ struct MockSingleOverload;
 // `mock_single_overload.gmock_Call(...)`. Because expectations are stored on
 // the MockingBitGen (an argument passed inside `Call(...)`), this forwards to
 // 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 "
                 "distribution result_type.");
-  using ArgTupleT = std::tuple<Args...>;
+  using KeyT = Ret(DistrT, std::tuple<Args...>);
   auto gmock_Call(
       absl::MockingBitGen& gen,  // NOLINT(google-runtime-references)
       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...);
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...)) {
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matchers...);
   }
 };
 
@@ -61,18 +58,14 @@ 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 "
                 "distribution result_type.");
-  using ArgTupleT = std::tuple<Arg, Args...>;
+  using KeyT = Ret(DistrT, std::tuple<Arg, Args...>);
   auto gmock_Call(
       const ::testing::Matcher<Arg>& matcher,
       absl::MockingBitGen& gen,  // NOLINT(google-runtime-references)
       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...);
+      -> decltype(MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher,
+                                                             matchers...)) {
+    return MockHelpers::MockFor<KeyT>(gen).gmock_Call(matcher, matchers...);
   }
 };