github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/xtest/waiters.go (about)

     1  package xtest
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    10  )
    11  
    12  const commonWaitTimeout = time.Second
    13  
    14  func WaitGroup(tb testing.TB, wg *sync.WaitGroup) {
    15  	tb.Helper()
    16  
    17  	WaitGroupWithTimeout(tb, wg, commonWaitTimeout)
    18  }
    19  
    20  func WaitGroupWithTimeout(tb testing.TB, wg *sync.WaitGroup, timeout time.Duration) {
    21  	tb.Helper()
    22  
    23  	groupFinished := make(empty.Chan)
    24  	go func() {
    25  		wg.Wait()
    26  		close(groupFinished)
    27  	}()
    28  
    29  	WaitChannelClosedWithTimeout(tb, groupFinished, timeout)
    30  }
    31  
    32  func WaitChannelClosed(t testing.TB, ch <-chan struct{}) {
    33  	t.Helper()
    34  
    35  	WaitChannelClosedWithTimeout(t, ch, commonWaitTimeout)
    36  }
    37  
    38  func WaitChannelClosedWithTimeout(t testing.TB, ch <-chan struct{}, timeout time.Duration) {
    39  	t.Helper()
    40  
    41  	select {
    42  	case <-time.After(timeout):
    43  		t.Fatal()
    44  	case <-ch:
    45  		// pass
    46  	}
    47  }
    48  
    49  // SpinWaitCondition wait while cond return true with check it in loop
    50  // l can be nil - then locker use for check conditions
    51  func SpinWaitCondition(tb testing.TB, l sync.Locker, cond func() bool) {
    52  	tb.Helper()
    53  
    54  	SpinWaitConditionWithTimeout(tb, l, commonWaitTimeout, cond)
    55  }
    56  
    57  // SpinWaitConditionWithTimeout wait while cond return true with check it in loop
    58  // l can be nil - then locker use for check conditions
    59  func SpinWaitConditionWithTimeout(tb testing.TB, l sync.Locker, condWaitTimeout time.Duration, cond func() bool) {
    60  	tb.Helper()
    61  
    62  	checkConditin := func() bool {
    63  		if l != nil {
    64  			l.Lock()
    65  			defer l.Unlock()
    66  		}
    67  
    68  		return cond()
    69  	}
    70  
    71  	start := time.Now()
    72  	for {
    73  		if checkConditin() {
    74  			return
    75  		}
    76  
    77  		if time.Since(start) > condWaitTimeout {
    78  			tb.Fatal()
    79  		}
    80  
    81  		runtime.Gosched()
    82  	}
    83  }
    84  
    85  // SpinWaitProgress failed if result of progress func no changes without timeout
    86  func SpinWaitProgress(tb testing.TB, progress func() (progressValue interface{}, finished bool)) {
    87  	tb.Helper()
    88  	SpinWaitProgressWithTimeout(tb, commonWaitTimeout, progress)
    89  }
    90  
    91  func SpinWaitProgressWithTimeout(
    92  	tb testing.TB,
    93  	timeout time.Duration,
    94  	progress func() (progressValue interface{}, finished bool),
    95  ) {
    96  	tb.Helper()
    97  
    98  	for {
    99  		currentValue, finished := progress()
   100  		if finished {
   101  			return
   102  		}
   103  
   104  		SpinWaitConditionWithTimeout(tb, nil, timeout, func() bool {
   105  			progressValue, finished := progress()
   106  
   107  			return finished || progressValue != currentValue
   108  		})
   109  	}
   110  }