Merge pull request go-kit#596 from go-kit/fix-test-flake

sd: fix TestDefaultEndpointer flake, hopefully

Author: peterbourgon

sd/endpointer_test.go

@@ -27,32 +27,42 @@ func TestDefaultEndpointer(t *testing.T) {
 	instancer.Update(sd.Event{Instances: []string{"a", "b"}})
 
 	endpointer := sd.NewEndpointer(instancer, f, log.NewNopLogger(), sd.InvalidateOnError(time.Minute))
-	if endpoints, err := endpointer.Endpoints(); err != nil {
-		t.Errorf("unepected error %v", err)
-	} else if want, have := 2, len(endpoints); want != have {
-		t.Errorf("want %d, have %d", want, have)
+
+	var (
+		endpoints []endpoint.Endpoint
+		err       error
+	)
+	if !within(time.Second, func() bool {
+		endpoints, err = endpointer.Endpoints()
+		return err == nil && len(endpoints) == 2
+	}) {
+		t.Errorf("wanted 2 endpoints, got %d (%v)", len(endpoints), err)
 	}
 
 	instancer.Update(sd.Event{Instances: []string{}})
+
 	select {
 	case <-ca:
 		t.Logf("endpoint a closed, good")
 	case <-time.After(time.Millisecond):
 		t.Errorf("didn't close the deleted instance in time")
 	}
+
 	select {
 	case <-cb:
 		t.Logf("endpoint b closed, good")
 	case <-time.After(time.Millisecond):
 		t.Errorf("didn't close the deleted instance in time")
 	}
+
 	if endpoints, err := endpointer.Endpoints(); err != nil {
 		t.Errorf("unepected error %v", err)
 	} else if want, have := 0, len(endpoints); want != have {
 		t.Errorf("want %d, have %d", want, have)
 	}
 
 	endpointer.Close()
+
 	instancer.Update(sd.Event{Instances: []string{"a"}})
 	// TODO verify that on Close the endpointer fully disconnects from the instancer.
 	// Unfortunately, because we use instance.Cache, this test cannot be in the sd package,
@@ -78,3 +88,14 @@ func (m *mockInstancer) Deregister(ch chan<- sd.Event) {
 type closer chan struct{}
 
 func (c closer) Close() error { close(c); return nil }
+
+func within(d time.Duration, f func() bool) bool {
+	deadline := time.Now().Add(d)
+	for time.Now().Before(deadline) {
+		if f() {
+			return true
+		}
+		time.Sleep(d / 10)
+	}
+	return false
+}

util/conn/manager_test.go

@@ -40,7 +40,7 @@ func TestManager(t *testing.T) {
 	// First takes should fail.
 	for i := 0; i < 10; i++ {
 		if conn = mgr.Take(); conn != nil {
-			t.Fatalf("want nil conn, got real conn")
+			t.Fatalf("iteration %d: want nil conn, got real conn", i)
 		}
 	}