Extend lntest/wait functions NoError and Predicate with function options

lntest/wait/wait.go

@@ -3,11 +3,41 @@
 import (
  "fmt"
  "time"
+
+ "github.com/lightningnetwork/lnd/fn"
 )
 
 // PollInterval is a constant specifying a 200 ms interval.
 const PollInterval = 200 * time.Millisecond
 
+// predicateOptions holds the options for the Predicate function.
+type predicateOptions struct {
+ // MaxCallAttempts is the maximum number of target function call
+ // attempts possible before returning an error.
+ MaxCallAttempts fn.Option[uint64]
+}
+
+// defaultPredicateOptions returns the default options for the Predicate
+// function.
+func defaultPredicateOptions() *predicateOptions {
+ return &predicateOptions{
+ MaxCallAttempts: fn.None[uint64](),
+ }
+}
+
+// PredicateOpt is a functional option that can be passed to the Predicate
+// function.
+type PredicateOpt func(*predicateOptions)
+
+// WithPredicateMaxCallAttempts is a functional option that can be passed to the
+// Predicate function to set the maximum number of target function call
+// attempts.
+func WithPredicateMaxCallAttempts(attempts uint64) PredicateOpt {
+ return func(o *predicateOptions) {
+ o.MaxCallAttempts = fn.Some(attempts)
+ }
+}
+
 // Predicate is a helper test function that will wait for a timeout period of
 // time until the passed predicate returns true. This function is helpful as
 // timing doesn't always line up well when running integration tests with
@@ -16,19 +46,33 @@
 //
 // TODO(yy): build a counter here so we know how many times we've tried the
 // `pred`.
-func Predicate(pred func() bool, timeout time.Duration) error {
+func Predicate(pred func() bool, timeout time.Duration,
+ opts ...PredicateOpt) error {
+
+ options := defaultPredicateOptions()
+ for _, opt := range opts {
+ opt(options)
+ }
+
  exitTimer := time.After(timeout)
  result := make(chan bool, 1)
 
+ // We'll keep track of the number of times we've called the predicate.
+ callAttemptsCounter := uint64(0)
+
  for {
  <-time.After(PollInterval)
 
+ // We'll increment the call attempts counter each time we call
+ // the predicate.
+ callAttemptsCounter += 1
+
  go func() {
  result <- pred()
  }()
 
  // Each time we call the pred(), we expect a result to be
- // returned otherwise it will timeout.
+ // returned otherwise it will time out.
  select {
  case <-exitTimer:
  return fmt.Errorf("predicate not satisfied after " +
@@ -38,14 +82,65 @@
  if succeed {
  return nil
  }
+
+ // If we have a max call attempts set, we'll check if
+ // we've exceeded it and return an error if so.
+ maxCallAttempts := options.MaxCallAttempts.UnwrapOr(
+ callAttemptsCounter + 1,
+ )
+ if callAttemptsCounter >= maxCallAttempts {
+ return fmt.Errorf("predicate not satisfied "+
+ "after max call attempts: %d",
+ maxCallAttempts)
+ }
  }
  }
 }
 
+// noErrorOptions holds the options for the NoError function.
+type noErrorOptions struct {
+ // maxCallAttempts is the maximum number of target function call
+ // attempts possible before returning an error.
+ maxCallAttempts fn.Option[uint64]
+}
+
+// defaultNoErrorOptions returns the default options for the NoError function.
+func defaultNoErrorOptions() *noErrorOptions {
+ return &noErrorOptions{
+ maxCallAttempts: fn.None[uint64](),
+ }
+}
+
+// NoErrorOpt is a functional option that can be passed to the NoError function.
+type NoErrorOpt func(*noErrorOptions)
+
+// WithNoErrorMaxCallAttempts is a functional option that can be passed to the
+// NoError function to set the maximum number of target function call attempts.
+func WithNoErrorMaxCallAttempts(attempts uint64) NoErrorOpt {
+ return func(o *noErrorOptions) {
+ o.maxCallAttempts = fn.Some(attempts)
+ }
+}
+
 // NoError is a wrapper around Predicate that waits for the passed method f to
 // execute without error, and returns the last error encountered if this doesn't
 // happen within the timeout.
-func NoError(f func() error, timeout time.Duration) error {
+func NoError(f func() error, timeout time.Duration,
+ opts ...NoErrorOpt) error {
+
+ options := defaultNoErrorOptions()
+ for _, opt := range opts {
+ opt(options)
+ }
+
+ // Formulate the options for the predicate function.
+ withPredicateMaxAttempts := func(*predicateOptions) {}
+ options.maxCallAttempts.WhenSome(func(attempts uint64) {
+ withPredicateMaxAttempts = WithPredicateMaxCallAttempts(
+ attempts,
+ )
+ })
+
  var predErr error
  pred := func() bool {
  if err := f(); err != nil {
@@ -57,7 +152,8 @@
 
  // If f() doesn't succeed within the timeout, return the last
  // encountered error.
- if err := Predicate(pred, timeout); err != nil {
+ err := Predicate(pred, timeout, withPredicateMaxAttempts)
+ if err != nil {
  // Handle the case where the passed in method, f, hangs for the
  // full timeout
  if predErr == nil {