go.uber.org/yarpc@v1.72.1/api/transport/transporttest/context.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package transporttest
    22  
    23  import (
    24  	"context"
    25  	"fmt"
    26  	"testing"
    27  	"time"
    28  )
    29  
    30  // ContextMatcher is a Matcher for verifying that a context's deadline is
    31  // within expected bounds: the current time, plus a TTL, plus or minus some
    32  // tolerance.
    33  type ContextMatcher struct {
    34  	t   *testing.T
    35  	ttl time.Duration
    36  
    37  	TTLDelta time.Duration
    38  }
    39  
    40  // ContextMatcherOption customizes the behavior of a ContextMatcher.
    41  type ContextMatcherOption interface {
    42  	run(*ContextMatcher)
    43  }
    44  
    45  // ContextTTL requires that a Context have the given TTL on it, with a
    46  // tolerance of TTLDelta.
    47  type ContextTTL time.Duration
    48  
    49  func (ttl ContextTTL) run(c *ContextMatcher) {
    50  	c.ttl = time.Duration(ttl)
    51  }
    52  
    53  // NewContextMatcher creates a ContextMatcher to verify properties about a
    54  // Context.
    55  func NewContextMatcher(t *testing.T, options ...ContextMatcherOption) *ContextMatcher {
    56  	matcher := &ContextMatcher{t: t, TTLDelta: DefaultTTLDelta}
    57  	for _, opt := range options {
    58  		opt.run(matcher)
    59  	}
    60  	return matcher
    61  }
    62  
    63  // Matches a context against an expected context, returning true only if the
    64  // given object is a context with a deadline that is now, plus the expected
    65  // TTL, plus or minus some tolerance.
    66  func (c *ContextMatcher) Matches(got interface{}) bool {
    67  	ctx, ok := got.(context.Context)
    68  	if !ok {
    69  		c.t.Logf("expected a Context but got a %T: %v", got, got)
    70  		return false
    71  	}
    72  
    73  	if c.ttl != 0 {
    74  		d, ok := ctx.Deadline()
    75  		if !ok {
    76  			c.t.Logf(
    77  				"expected Context to have a TTL of %v but it has no deadline", c.ttl)
    78  			return false
    79  		}
    80  
    81  		ttl := time.Until(d)
    82  		maxTTL := c.ttl + c.TTLDelta
    83  		minTTL := c.ttl - c.TTLDelta
    84  		if ttl > maxTTL || ttl < minTTL {
    85  			c.t.Logf("TTL out of expected bounds: %v < %v < %v", minTTL, ttl, maxTTL)
    86  			return false
    87  		}
    88  	}
    89  
    90  	return true
    91  }
    92  
    93  func (c *ContextMatcher) String() string {
    94  	return fmt.Sprintf("ContextMatcher(TTL:%v±%v)", c.ttl, c.TTLDelta)
    95  }