test(responsemanager): fix tests

responsemanager/responsemanager_test.go

@@ -222,9 +222,18 @@ func TestCancellationQueryInProgress(t *testing.T) {
 
 	responseManager.synchronize()
 
-	// read one block -- to unblock processing
+	// at this point we should receive at most one more block, then traversal
+	// should complete
+	additionalMessageCount := 0
+drainqueue:
+	for {
 		select {
+		case <-ctx.Done():
+			t.Fatal("Should have completed request but didn't")
 		case sentResponse := <-sentResponses:
+			if additionalMessageCount > 0 {
+				t.Fatal("should not send any more responses")
+			}
 			k := sentResponse.link.(cidlink.Link)
 			blockIndex := testutil.IndexOf(blks, k.Cid)
 			if blockIndex == -1 {
@@ -236,17 +245,10 @@ func TestCancellationQueryInProgress(t *testing.T) {
 			if sentResponse.requestID != requestID {
 				t.Fatal("incorrect response id")
 			}
-	case <-ctx.Done():
-		t.Fatal("did not send responses")
-	}
-
-	// at this point traversal should abort and we should receive a completion
-	select {
-	case <-ctx.Done():
-		t.Fatal("Should have completed request but didn't")
-	case <-sentResponses:
-		t.Fatal("should not send any more responses")
+			additionalMessageCount++
 		case <-requestIDChan:
+			break drainqueue
+		}
 	}
 }