forked from DataDog/datadog-agent
-
Notifications
You must be signed in to change notification settings - Fork 2
/
timeout.go
72 lines (64 loc) · 1.7 KB
/
timeout.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
package testutil
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
// AssertTrueBeforeTimeout regularly checks whether a condition is met. It
// does so until a timeout is reached, in which case it makes the test fail.
// Condition is evaluated in a goroutine to avoid tests hanging if a system
// is deadlocked.
func AssertTrueBeforeTimeout(t *testing.T, frequency, timeout time.Duration, condition func() bool) {
t.Helper()
assertTrueBeforeTimeout(t, frequency, timeout, condition)
}
// RequireTrueBeforeTimeout is the same as AssertTrueBeforeTimeout, but it calls
// t.failNow() if the condition function times out.
func RequireTrueBeforeTimeout(t *testing.T, frequency, timeout time.Duration, condition func() bool) {
t.Helper()
if !assertTrueBeforeTimeout(t, frequency, timeout, condition) {
t.FailNow()
}
}
func assertTrueBeforeTimeout(t *testing.T, frequency, timeout time.Duration, condition func() bool) bool {
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
r := make(chan bool, 1)
go func() {
// Try once immediately
r <- condition()
// Retry until timeout
checkTicker := time.NewTicker(frequency)
defer checkTicker.Stop()
for {
select {
case <-checkTicker.C:
ok := condition()
r <- ok
if ok {
return
}
case <-ctx.Done():
return
}
}
}()
var ranOnce bool
for {
select {
case ok := <-r:
if ok {
return ok
}
ranOnce = true
case <-ctx.Done():
if ranOnce {
assert.Fail(t, "Timeout waiting for condition to happen, function returned false")
} else {
assert.Fail(t, "Timeout waiting for condition to happen, function never returned")
}
return false
}
}
}