github.com/yandex/pandora@v0.5.32/core/coretest/schedule.go (about)

     1  package coretest
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/require"
     8  	"github.com/yandex/pandora/core"
     9  )
    10  
    11  func ExpectScheduleNextsStartAt(t *testing.T, sched core.Schedule, startAt time.Time, nexts ...time.Duration) {
    12  	beforeStartLeft := sched.Left()
    13  	tokensExpected := len(nexts) - 1 // Last next is finish time.
    14  	require.Equal(t, tokensExpected, beforeStartLeft)
    15  	sched.Start(startAt)
    16  	actualNexts := DrainScheduleDuration(t, sched, startAt)
    17  	require.Equal(t, nexts, actualNexts)
    18  }
    19  
    20  func ExpectScheduleNexts(t *testing.T, sched core.Schedule, nexts ...time.Duration) {
    21  	ExpectScheduleNextsStartAt(t, sched, time.Now(), nexts...)
    22  }
    23  
    24  func ExpectScheduleNextsStartAtT(t *testing.T, sched core.Schedule, startAt time.Time, nexts ...time.Duration) {
    25  	beforeStartLeft := sched.Left()
    26  	tokensExpected := len(nexts) - 1 // Last next is finish time.
    27  	require.Equal(t, tokensExpected, beforeStartLeft)
    28  	sched.Start(startAt)
    29  	actualNexts := DrainScheduleDuration(t, sched, startAt)
    30  	require.Equal(t, nexts, actualNexts)
    31  }
    32  
    33  func ExpectScheduleNextsT(t *testing.T, sched core.Schedule, nexts ...time.Duration) {
    34  	ExpectScheduleNextsStartAtT(t, sched, time.Now(), nexts...)
    35  }
    36  
    37  const drainLimit = 1000000
    38  
    39  // DrainSchedule starts schedule and takes all tokens from it.
    40  // Returns all tokens and finish time relative to start
    41  func DrainScheduleDuration(t *testing.T, sched core.Schedule, startAt time.Time) []time.Duration {
    42  	nexts := DrainSchedule(t, sched)
    43  	durations := make([]time.Duration, len(nexts))
    44  	for i, next := range nexts {
    45  		durations[i] = next.Sub(startAt)
    46  	}
    47  	return durations
    48  }
    49  
    50  // DrainSchedule takes all tokens from passed schedule.
    51  // Returns all tokens and finish time.
    52  func DrainSchedule(t *testing.T, sched core.Schedule) []time.Time {
    53  	expectedLeft := sched.Left()
    54  	var nexts []time.Time
    55  	for len(nexts) < drainLimit {
    56  		next, ok := sched.Next()
    57  		nexts = append(nexts, next)
    58  		if !ok {
    59  			require.Equal(t, 0, sched.Left())
    60  			return nexts
    61  		}
    62  		expectedLeft--
    63  		require.Equal(t, expectedLeft, sched.Left())
    64  	}
    65  	panic("drain limit reached")
    66  }