fix a flaky test
diff --git a/picker_test.go b/picker_test.go
index efe0c23..70dae20 100644
--- a/picker_test.go
+++ b/picker_test.go
@@ -36,6 +36,7 @@
import (
"fmt"
"math"
+ "sync"
"testing"
"time"
@@ -50,6 +51,7 @@
side chan int
// the channel to notifiy update injector that the update reading is done
readDone chan int
+ wg *sync.WaitGroup
}
func (w *testWatcher) Next() (updates []*naming.Update, err error) {
@@ -81,6 +83,7 @@
type testNameResolver struct {
w *testWatcher
addr string
+ wg *sync.WaitGroup
}
func (r *testNameResolver) Resolve(target string) (naming.Watcher, error) {
@@ -88,6 +91,7 @@
update: make(chan *naming.Update, 1),
side: make(chan int, 1),
readDone: make(chan int),
+ wg: r.wg,
}
r.w.side <- 1
r.w.update <- &naming.Update{
@@ -96,11 +100,14 @@
}
go func() {
<-r.w.readDone
+ if r.w.wg != nil {
+ r.w.wg.Done()
+ }
}()
return r.w, nil
}
-func startServers(t *testing.T, numServers, port int, maxStreams uint32) ([]*server, *testNameResolver) {
+func startServers(t *testing.T, numServers, port int, maxStreams uint32, wg *sync.WaitGroup) ([]*server, *testNameResolver) {
var servers []*server
for i := 0; i < numServers; i++ {
s := &server{readyChan: make(chan bool)}
@@ -110,12 +117,15 @@
}
// Point to server1
addr := "127.0.0.1:" + servers[0].port
- return servers, &testNameResolver{addr: addr}
+ return servers, &testNameResolver{
+ addr: addr,
+ wg: wg,
+ }
}
func TestNameDiscovery(t *testing.T) {
// Start 3 servers on 3 ports.
- servers, r := startServers(t, 3, 0, math.MaxUint32)
+ servers, r := startServers(t, 3, 0, math.MaxUint32, nil)
cc, err := Dial("foo.bar.com", WithPicker(NewUnicastNamingPicker(r)), WithBlock(), WithInsecure(), WithCodec(testCodec{}))
if err != nil {
t.Fatalf("Failed to create ClientConn: %v", err)
@@ -156,7 +166,9 @@
}
func TestEmptyAddrs(t *testing.T) {
- servers, r := startServers(t, 1, 0, math.MaxUint32)
+ var wg sync.WaitGroup
+ servers, r := startServers(t, 1, 0, math.MaxUint32, &wg)
+ wg.Add(1)
cc, err := Dial("foo.bar.com", WithPicker(NewUnicastNamingPicker(r)), WithBlock(), WithInsecure(), WithCodec(testCodec{}))
if err != nil {
t.Fatalf("Failed to create ClientConn: %v", err)
@@ -172,6 +184,7 @@
Op: naming.Delete,
Addr: "127.0.0.1:" + servers[0].port,
})
+ wg.Wait()
r.w.inject(updates)
ctx, _ := context.WithTimeout(context.Background(), 10*time.Millisecond)
if err := Invoke(ctx, "/foo/bar", &expectedRequest, &reply, cc); err == nil {
