...
1
18
19 package testutils
20
21 import (
22 "context"
23 "testing"
24
25 "google.golang.org/grpc/connectivity"
26 )
27
28
29 type StateChanger interface {
30
31 Connect()
32
33 GetState() connectivity.State
34
35
36 WaitForStateChange(ctx context.Context, s connectivity.State) bool
37 }
38
39
40
41 func StayConnected(ctx context.Context, sc StateChanger) {
42 for {
43 state := sc.GetState()
44 switch state {
45 case connectivity.Idle:
46 sc.Connect()
47 case connectivity.Shutdown:
48 return
49 }
50 if !sc.WaitForStateChange(ctx, state) {
51 return
52 }
53 }
54 }
55
56
57
58 func AwaitState(ctx context.Context, t *testing.T, sc StateChanger, stateWant connectivity.State) {
59 t.Helper()
60 for state := sc.GetState(); state != stateWant; state = sc.GetState() {
61 if !sc.WaitForStateChange(ctx, state) {
62 t.Fatalf("Timed out waiting for state change. got %v; want %v", state, stateWant)
63 }
64 }
65 }
66
67
68
69 func AwaitNotState(ctx context.Context, t *testing.T, sc StateChanger, stateDoNotWant connectivity.State) {
70 t.Helper()
71 for state := sc.GetState(); state == stateDoNotWant; state = sc.GetState() {
72 if !sc.WaitForStateChange(ctx, state) {
73 t.Fatalf("Timed out waiting for state change. got %v; want NOT %v", state, stateDoNotWant)
74 }
75 }
76 }
77
78
79
80 func AwaitNoStateChange(ctx context.Context, t *testing.T, sc StateChanger, currState connectivity.State) {
81 t.Helper()
82 if sc.WaitForStateChange(ctx, currState) {
83 t.Fatalf("State changed from %q to %q when no state change was expected", currState, sc.GetState())
84 }
85 }
86
View as plain text