github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/internal/testingutil/poll.go (about)

     1  // Copyright 2018 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package testingutil
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"time"
    11  
    12  	"go.chromium.org/tast/core/ctxutil"
    13  	"go.chromium.org/tast/core/errors"
    14  )
    15  
    16  const defaultPollInterval = 100 * time.Millisecond
    17  
    18  // PollOptions provides testing.PollOptions.
    19  type PollOptions struct {
    20  	// Timeout specifies the maximum time to poll.
    21  	// Non-positive values indicate no timeout (although context deadlines will still be honored).
    22  	Timeout time.Duration
    23  	// Interval specifies how long to sleep between polling.
    24  	// Non-positive values indicate that a reasonable default should be used.
    25  	Interval time.Duration
    26  }
    27  
    28  // pollBreak is a wrapper of error to terminate the Poll immediately.
    29  type pollBreak struct {
    30  	err error
    31  }
    32  
    33  // Error implementation of pollBreak. However, it is not expected that this
    34  // is used directly, since pollBreak is not returned to callers.
    35  func (b *pollBreak) Error() string {
    36  	return b.err.Error()
    37  }
    38  
    39  // PollBreak implements testing.PollBreak.
    40  func PollBreak(err error) error {
    41  	return &pollBreak{err}
    42  }
    43  
    44  // Poll implements testing.Poll.
    45  func Poll(ctx context.Context, f func(context.Context) error, opts *PollOptions) error {
    46  	if ctx.Err() != nil {
    47  		return errors.Wrap(ctx.Err(), "poll fails before actually running the function")
    48  	}
    49  
    50  	timeout := ctxutil.MaxTimeout
    51  	timeoutLog := "with no set timeout"
    52  	if opts != nil && opts.Timeout > 0 {
    53  		timeout = opts.Timeout
    54  		timeoutLog = fmt.Sprintf("with timeout %v", timeout)
    55  	}
    56  	ctx, cancel := context.WithTimeout(ctx, timeout)
    57  	defer cancel()
    58  
    59  	interval := defaultPollInterval
    60  	if opts != nil && opts.Interval > 0 {
    61  		interval = opts.Interval
    62  	}
    63  
    64  	var lastErr error
    65  	for {
    66  		var err error
    67  		if err = f(ctx); err == nil {
    68  			return nil
    69  		}
    70  
    71  		if e, ok := err.(*pollBreak); ok {
    72  			if ctx.Err() != nil && lastErr != nil {
    73  				return errors.Wrapf(lastErr, "%s during a poll %v; last error follows", e.err, timeoutLog)
    74  			}
    75  			return e.err
    76  		}
    77  
    78  		// If f honors ctx's deadline, it may return a "context deadline exceeded" error
    79  		// if the deadline is reached while is running. To avoid returning a useless
    80  		// "context deadline exceeded; last error follows: context deadline exceeded)" error below,
    81  		// save the last error that is returned before the deadline is reached.
    82  		if lastErr == nil || ctx.Err() == nil {
    83  			lastErr = err
    84  		}
    85  
    86  		select {
    87  		case <-time.After(interval):
    88  		case <-ctx.Done():
    89  			if lastErr != nil {
    90  				return errors.Wrapf(lastErr, "%s during a poll %v; last error follows", ctx.Err(), timeoutLog)
    91  			}
    92  			return errors.Wrap(ctx.Err(), "poll fails before the first execution of the given function completes")
    93  		}
    94  	}
    95  }