1 files changed, 33 insertions, 7 deletions
diff --git a/absl/random/internal/mock_helpers.h b/absl/random/internal/mock_helpers.h
index a7a97bfc..19d05612 100644
--- a/absl/random/internal/mock_helpers.h
+++ b/absl/random/internal/mock_helpers.h
@@ -16,10 +16,9 @@
 #ifndef ABSL_RANDOM_INTERNAL_MOCK_HELPERS_H_
 #define ABSL_RANDOM_INTERNAL_MOCK_HELPERS_H_
 
-#include <tuple>
-#include <type_traits>
 #include <utility>
 
+#include "absl/base/config.h"
 #include "absl/base/internal/fast_type_id.h"
 #include "absl/types/optional.h"
 
@@ -27,6 +26,16 @@ namespace absl {
 ABSL_NAMESPACE_BEGIN
 namespace random_internal {
 
+// A no-op validator meeting the ValidatorT requirements for MockHelpers.
+//
+// Custom validators should follow a similar structure, passing the type to
+// MockHelpers::MockFor<KeyT>(m, CustomValidatorT()).
+struct NoOpValidator {
+  // Default validation: do nothing.
+  template <typename ResultT, typename... Args>
+  static void Validate(ResultT, Args&&...) {}
+};
+
 // MockHelpers works in conjunction with MockOverloadSet, MockingBitGen, and
 // BitGenRef to enable the mocking capability for absl distribution functions.
 //
@@ -109,22 +118,39 @@ class MockHelpers {
         0, urbg, std::forward<Args>(args)...);
   }
 
-  // Acquire a mock for the KeyT (may or may not be a signature).
+  // Acquire a mock for the KeyT (may or may not be a signature), set up to use
+  // the ValidatorT to verify that the result is in the range of the RNG
+  // function.
   //
   // KeyT is used to generate a typeid-based lookup for the mock.
   // KeyT is a signature of the form:
   //   result_type(discriminator_type, std::tuple<args...>)
   // The mocked function signature will be composed from KeyT as:
   //   result_type(args...)
-  template <typename KeyT, typename MockURBG>
-  static auto MockFor(MockURBG& m)
+  // ValidatorT::Validate will be called after the result of the RNG. The
+  //   signature is expected to be of the form:
+  //      ValidatorT::Validate(result, args...)
+  template <typename KeyT, typename ValidatorT, typename MockURBG>
+  static auto MockFor(MockURBG& m, ValidatorT)
       -> decltype(m.template RegisterMock<
                   typename KeySignature<KeyT>::result_type,
                   typename KeySignature<KeyT>::arg_tuple_type>(
-          m, std::declval<IdType>())) {
+          m, std::declval<IdType>(), ValidatorT())) {
     return m.template RegisterMock<typename KeySignature<KeyT>::result_type,
                                    typename KeySignature<KeyT>::arg_tuple_type>(
-        m, ::absl::base_internal::FastTypeId<KeyT>());
+        m, ::absl::base_internal::FastTypeId<KeyT>(), ValidatorT());
+  }
+
+  // Acquire a mock for the KeyT (may or may not be a signature).
+  //
+  // KeyT is used to generate a typeid-based lookup for the mock.
+  // KeyT is a signature of the form:
+  //   result_type(discriminator_type, std::tuple<args...>)
+  // The mocked function signature will be composed from KeyT as:
+  //   result_type(args...)
+  template <typename KeyT, typename MockURBG>
+  static decltype(auto) MockFor(MockURBG& m) {
+    return MockFor<KeyT>(m, NoOpValidator());
   }
 };