istio.io/istio@v0.0.0-20240520182934-d79c90f27776/pkg/test/util/assert/tracker.go (about)

     1  // Copyright Istio 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 assert
    16  
    17  import (
    18  	"fmt"
    19  	"sync"
    20  	"time"
    21  
    22  	"istio.io/istio/pkg/ptr"
    23  	"istio.io/istio/pkg/test"
    24  	"istio.io/istio/pkg/test/util/retry"
    25  )
    26  
    27  type Tracker[T comparable] struct {
    28  	t      test.Failer
    29  	mu     sync.Mutex
    30  	events []T
    31  }
    32  
    33  // NewTracker builds a tracker which records events that occur
    34  func NewTracker[T comparable](t test.Failer) *Tracker[T] {
    35  	return &Tracker[T]{t: t}
    36  }
    37  
    38  // Record that an event occurred.
    39  func (t *Tracker[T]) Record(event T) {
    40  	t.mu.Lock()
    41  	defer t.mu.Unlock()
    42  	t.events = append(t.events, event)
    43  }
    44  
    45  // Empty asserts the tracker is empty
    46  func (t *Tracker[T]) Empty() {
    47  	t.t.Helper()
    48  	t.mu.Lock()
    49  	defer t.mu.Unlock()
    50  	if len(t.events) != 0 {
    51  		t.t.Fatalf("unexpected events: %v", t.events)
    52  	}
    53  }
    54  
    55  // WaitOrdered waits for an event to happen, in order
    56  func (t *Tracker[T]) WaitOrdered(events ...T) {
    57  	t.t.Helper()
    58  	for i, event := range events {
    59  		var err error
    60  		retry.UntilSuccessOrFail(t.t, func() error {
    61  			t.mu.Lock()
    62  			defer t.mu.Unlock()
    63  			if len(t.events) == 0 {
    64  				return fmt.Errorf("no events")
    65  			}
    66  			if t.events[0] != event {
    67  				// Exit early instead of continuing to retry
    68  				err = fmt.Errorf("got events %v, want %v (%d)", t.events, event, i)
    69  				return nil
    70  			}
    71  			// clear the event
    72  			t.events[0] = ptr.Empty[T]()
    73  			t.events = t.events[1:]
    74  			return nil
    75  		}, retry.Timeout(time.Second), retry.BackoffDelay(time.Millisecond))
    76  		if err != nil {
    77  			t.t.Fatal(err)
    78  		}
    79  	}
    80  	t.Empty()
    81  }
    82  
    83  func (t *Tracker[T]) Events() []T {
    84  	t.mu.Lock()
    85  	defer t.mu.Unlock()
    86  	return t.events
    87  }
    88  
    89  // WaitUnordered waits for an event to happen, in any order
    90  func (t *Tracker[T]) WaitUnordered(events ...T) {
    91  	t.t.Helper()
    92  	want := map[T]struct{}{}
    93  	for _, e := range events {
    94  		want[e] = struct{}{}
    95  	}
    96  	var err error
    97  	retry.UntilSuccessOrFail(t.t, func() error {
    98  		t.mu.Lock()
    99  		defer t.mu.Unlock()
   100  		if len(t.events) == 0 {
   101  			return fmt.Errorf("no events (want %v)", want)
   102  		}
   103  		got := t.events[0]
   104  		if _, f := want[got]; !f {
   105  			// Exit early instead of continuing to retry
   106  			err = fmt.Errorf("got events %v, want %v", t.events, want)
   107  			return nil
   108  		}
   109  		// clear the event
   110  		t.events[0] = ptr.Empty[T]()
   111  		t.events = t.events[1:]
   112  		delete(want, got)
   113  
   114  		if len(want) > 0 {
   115  			return fmt.Errorf("still waiting for %v", want)
   116  		}
   117  		return nil
   118  	}, retry.Timeout(time.Second), retry.BackoffDelay(time.Millisecond))
   119  	if err != nil {
   120  		t.t.Fatal(err)
   121  	}
   122  	t.Empty()
   123  }
   124  
   125  // WaitCompare waits for an event to happen and ensures it meets a custom comparison function
   126  func (t *Tracker[T]) WaitCompare(f func(T) bool) {
   127  	t.t.Helper()
   128  	var err error
   129  	retry.UntilSuccessOrFail(t.t, func() error {
   130  		t.mu.Lock()
   131  		defer t.mu.Unlock()
   132  		if len(t.events) == 0 {
   133  			return fmt.Errorf("no events")
   134  		}
   135  		got := t.events[0]
   136  		if !f(got) {
   137  			// Exit early instead of continuing to retry
   138  			err = fmt.Errorf("got events %v, which does not match criteria", t.events)
   139  			return nil
   140  		}
   141  		// clear the event
   142  		t.events[0] = ptr.Empty[T]()
   143  		t.events = t.events[1:]
   144  		return nil
   145  	}, retry.Timeout(time.Second), retry.BackoffDelay(time.Millisecond))
   146  	if err != nil {
   147  		t.t.Fatal(err)
   148  	}
   149  	t.Empty()
   150  }