GODRIVER-2667 PR Changes

Author: timothy-kim-mongo

mongo/integration/unified_runner_events_helper_test.go

@@ -88,6 +88,13 @@ func waitForEvent(mt *mtest.T, test *testCase, op *operation) {
 	expectedCount := int(op.Arguments.Lookup("count").Int32())
 
 	callback := func() bool {
+		// Stop loop if callback has been canceled.
+		select {
+		case <-context.Background().Done():
+			return true
+		default:
+		}
+
 		var count int
 		// Spec tests only ever wait for ServerMarkedUnknown SDAM events for the time being.
 		if eventType == "ServerMarkedUnknownEvent" {
@@ -126,11 +133,20 @@ func recordPrimary(mt *mtest.T, testCase *testCase) {
 }
 
 func waitForPrimaryChange(mt *mtest.T, testCase *testCase, op *operation) {
+	callback := func() bool {
+		// Stop loop if callback has been canceled.
+		select {
+		case <-context.Background().Done():
+			return true
+		default:
+		}
+
+		return getPrimaryAddress(mt, testCase.testTopology, false) != testCase.recordedPrimary
+	}
+
 	timeout := convertValueToMilliseconds(mt, op.Arguments.Lookup("timeoutMS"))
 	assert.Eventually(mt,
-		func() bool {
-			return getPrimaryAddress(mt, testCase.testTopology, false) != testCase.recordedPrimary
-		},
+		callback,
 		timeout,
 		100*time.Millisecond,
 		"expected primary address to be different within the timeout period")