go.etcd.io/etcd@v3.3.27+incompatible/pkg/testutil/recorder.go (about)

     1  // Copyright 2015 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package testutil
    16  
    17  import (
    18  	"errors"
    19  	"fmt"
    20  	"sync"
    21  	"time"
    22  )
    23  
    24  type Action struct {
    25  	Name   string
    26  	Params []interface{}
    27  }
    28  
    29  type Recorder interface {
    30  	// Record publishes an Action (e.g., function call) which will
    31  	// be reflected by Wait() or Chan()
    32  	Record(a Action)
    33  	// Wait waits until at least n Actions are available or returns with error
    34  	Wait(n int) ([]Action, error)
    35  	// Action returns immediately available Actions
    36  	Action() []Action
    37  	// Chan returns the channel for actions published by Record
    38  	Chan() <-chan Action
    39  }
    40  
    41  // RecorderBuffered appends all Actions to a slice
    42  type RecorderBuffered struct {
    43  	sync.Mutex
    44  	actions []Action
    45  }
    46  
    47  func (r *RecorderBuffered) Record(a Action) {
    48  	r.Lock()
    49  	r.actions = append(r.actions, a)
    50  	r.Unlock()
    51  }
    52  func (r *RecorderBuffered) Action() []Action {
    53  	r.Lock()
    54  	cpy := make([]Action, len(r.actions))
    55  	copy(cpy, r.actions)
    56  	r.Unlock()
    57  	return cpy
    58  }
    59  func (r *RecorderBuffered) Wait(n int) (acts []Action, err error) {
    60  	// legacy racey behavior
    61  	WaitSchedule()
    62  	acts = r.Action()
    63  	if len(acts) < n {
    64  		err = newLenErr(n, len(acts))
    65  	}
    66  	return acts, err
    67  }
    68  
    69  func (r *RecorderBuffered) Chan() <-chan Action {
    70  	ch := make(chan Action)
    71  	go func() {
    72  		acts := r.Action()
    73  		for i := range acts {
    74  			ch <- acts[i]
    75  		}
    76  		close(ch)
    77  	}()
    78  	return ch
    79  }
    80  
    81  // RecorderStream writes all Actions to an unbuffered channel
    82  type recorderStream struct {
    83  	ch chan Action
    84  }
    85  
    86  func NewRecorderStream() Recorder {
    87  	return &recorderStream{ch: make(chan Action)}
    88  }
    89  
    90  func (r *recorderStream) Record(a Action) {
    91  	r.ch <- a
    92  }
    93  
    94  func (r *recorderStream) Action() (acts []Action) {
    95  	for {
    96  		select {
    97  		case act := <-r.ch:
    98  			acts = append(acts, act)
    99  		default:
   100  			return acts
   101  		}
   102  	}
   103  }
   104  
   105  func (r *recorderStream) Chan() <-chan Action {
   106  	return r.ch
   107  }
   108  
   109  func (r *recorderStream) Wait(n int) ([]Action, error) {
   110  	acts := make([]Action, n)
   111  	timeoutC := time.After(5 * time.Second)
   112  	for i := 0; i < n; i++ {
   113  		select {
   114  		case acts[i] = <-r.ch:
   115  		case <-timeoutC:
   116  			acts = acts[:i]
   117  			return acts, newLenErr(n, i)
   118  		}
   119  	}
   120  	// extra wait to catch any Action spew
   121  	select {
   122  	case act := <-r.ch:
   123  		acts = append(acts, act)
   124  	case <-time.After(10 * time.Millisecond):
   125  	}
   126  	return acts, nil
   127  }
   128  
   129  func newLenErr(expected int, actual int) error {
   130  	s := fmt.Sprintf("len(actions) = %d, expected >= %d", actual, expected)
   131  	return errors.New(s)
   132  }