github.com/searKing/golang/go@v1.2.117/testing/channel.go (about)

     1  // Copyright 2021 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package testing
     6  
     7  import (
     8  	"context"
     9  )
    10  
    11  // DefaultChanBufferSize is the default buffer size of the underlying channel.
    12  const DefaultChanBufferSize = 1
    13  
    14  // Channel wraps a generic channel and provides a timed receive operation.
    15  type Channel struct {
    16  	ch chan any
    17  }
    18  
    19  // Send sends value on the underlying channel.
    20  func (c *Channel) Send(value any) {
    21  	c.ch <- value
    22  }
    23  
    24  // SendContext sends value on the underlying channel, or returns an error if
    25  // the context expires.
    26  func (c *Channel) SendContext(ctx context.Context, value any) error {
    27  	select {
    28  	case c.ch <- value:
    29  		return nil
    30  	case <-ctx.Done():
    31  		return ctx.Err()
    32  	}
    33  }
    34  
    35  // SendOrFail attempts to send value on the underlying channel.  Returns true
    36  // if successful or false if the channel was full.
    37  func (c *Channel) SendOrFail(value any) bool {
    38  	select {
    39  	case c.ch <- value:
    40  		return true
    41  	default:
    42  		return false
    43  	}
    44  }
    45  
    46  // ReceiveOrFail returns the value on the underlying channel and true, or nil
    47  // and false if the channel was empty.
    48  func (c *Channel) ReceiveOrFail() (any, bool) {
    49  	select {
    50  	case got := <-c.ch:
    51  		return got, true
    52  	default:
    53  		return nil, false
    54  	}
    55  }
    56  
    57  // Receive returns the value received on the underlying channel, or the error
    58  // returned by ctx if it is closed or cancelled.
    59  func (c *Channel) Receive(ctx context.Context) (any, error) {
    60  	select {
    61  	case <-ctx.Done():
    62  		return nil, ctx.Err()
    63  	case got := <-c.ch:
    64  		return got, nil
    65  	}
    66  }
    67  
    68  // Replace clears the value on the underlying channel, and sends the new value.
    69  //
    70  // It's expected to be used with a size-1 channel, to only keep the most
    71  // up-to-date item. This method is inherently racy when invoked concurrently
    72  // from multiple goroutines.
    73  func (c *Channel) Replace(value any) {
    74  	for {
    75  		select {
    76  		case c.ch <- value:
    77  			return
    78  		case <-c.ch:
    79  		}
    80  	}
    81  }
    82  
    83  // Clear drains all values on the underlying channel, or the error
    84  // returned by ctx if it is closed or cancelled.
    85  func (c *Channel) Clear(ctx context.Context) error {
    86  	for {
    87  		select {
    88  		case <-ctx.Done():
    89  			return ctx.Err()
    90  		case <-c.ch:
    91  			continue
    92  		default:
    93  			return nil
    94  		}
    95  	}
    96  }
    97  
    98  // NewChannel returns a new Channel.
    99  func NewChannel() *Channel {
   100  	return NewChannelWithSize(DefaultChanBufferSize)
   101  }
   102  
   103  // NewChannelWithSize returns a new Channel with a buffer of bufSize.
   104  func NewChannelWithSize(bufSize int) *Channel {
   105  	return &Channel{ch: make(chan any, bufSize)}
   106  }