From f3e35d307deff977ac05d4347635d53ea3400be3 Mon Sep 17 00:00:00 2001
From: hannahhoward <hannah@hannahhoward.net>
Date: Thu, 7 Oct 2021 16:14:03 -0700
Subject: [PATCH] test(responsemanager): fix flaky
 TestCancellationQueryInProgress

resolve possible failure from request finishing early -- simply to fix for
TestCancellationViaCommand
---
 responsemanager/responsemanager_test.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/responsemanager/responsemanager_test.go b/responsemanager/responsemanager_test.go
index cfe1365e..937bcea7 100644
--- a/responsemanager/responsemanager_test.go
+++ b/responsemanager/responsemanager_test.go
@@ -73,6 +73,15 @@ func TestCancellationQueryInProgress(t *testing.T) {
 	defer td.cancel()
 	responseManager := td.newResponseManager()
 	td.requestHooks.Register(selectorvalidator.SelectorValidator(100))
+	// This block hook is simply used to pause block hook processing after 1 block until the cancel command is sent
+	blkCount := 0
+	waitForCancel := make(chan struct{})
+	td.blockHooks.Register(func(p peer.ID, requestData graphsync.RequestData, blockData graphsync.BlockData, hookActions graphsync.OutgoingBlockHookActions) {
+		if blkCount == 1 {
+			<-waitForCancel
+		}
+		blkCount++
+	})
 	cancelledListenerCalled := make(chan struct{}, 1)
 	td.cancelledListeners.Register(func(p peer.ID, request graphsync.RequestData) {
 		td.connManager.RefuteProtected(t, td.p)
@@ -90,6 +99,7 @@ func TestCancellationQueryInProgress(t *testing.T) {
 	}
 	responseManager.ProcessRequests(td.ctx, td.p, cancelRequests)
 	responseManager.synchronize()
+	close(waitForCancel)
 
 	testutil.AssertDoesReceive(td.ctx, t, cancelledListenerCalled, "should call cancelled listener")