integration: don't expect recv to stop on CloseSend in waitResponse

This commit is contained in:
Anthony Romano 2016-11-18 11:36:24 -08:00
parent 456569f45d
commit 6d9168a2ec

View File

@ -948,21 +948,23 @@ func testV3WatchMultipleStreams(t *testing.T, startRev int64) {
// returned closing the WatchClient stream. Or the response will // returned closing the WatchClient stream. Or the response will
// be returned. // be returned.
func waitResponse(wc pb.Watch_WatchClient, timeout time.Duration) (bool, *pb.WatchResponse) { func waitResponse(wc pb.Watch_WatchClient, timeout time.Duration) (bool, *pb.WatchResponse) {
rCh := make(chan *pb.WatchResponse) rCh := make(chan *pb.WatchResponse, 1)
donec := make(chan struct{})
defer close(donec)
go func() { go func() {
resp, _ := wc.Recv() resp, _ := wc.Recv()
rCh <- resp select {
case rCh <- resp:
case <-donec:
}
}() }()
select { select {
case nr := <-rCh: case nr := <-rCh:
return false, nr return false, nr
case <-time.After(timeout): case <-time.After(timeout):
} }
// didn't get response
wc.CloseSend() wc.CloseSend()
rv, ok := <-rCh
if rv != nil || !ok {
return false, rv
}
return true, nil return true, nil
} }