Merge pull request #420 from twmb/test_fixup

kgo tests: improve resilience

pkg/kgo/consumer_direct_test.go

@@ -21,11 +21,11 @@ func TestIssue325(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 	cl.AddConsumeTopics(topic)
 	recs := cl.PollFetches(ctx).Records()
-	if len(recs) != 1 && string(recs[0].Value) != "foo" {
+	if len(recs) != 1 || string(recs[0].Value) != "foo" {
 		t.Fatal(recs)
 	}
 }
@@ -53,7 +53,7 @@ func TestIssue337(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	var recs []*Record
 out:

pkg/kgo/group_test.go

@@ -125,6 +125,7 @@ func (c *testConsumer) etl(etlsBeforeQuit int) {
 		Balancers(c.balancer),
 		MaxBufferedRecords(10000),
 		ConsumePreferringLagFn(PreferLagAt(1)),
+		BlockRebalanceOnPoll(),
 
 		// Even with autocommitting, autocommitting does not commit
 		// *the latest* when being revoked. We always want to commit
@@ -160,7 +161,9 @@ func (c *testConsumer) etl(etlsBeforeQuit int) {
 		}
 	}()
 
+	defer cl.AllowRebalance()
 	for {
+		cl.AllowRebalance()
 		// If we etl a few times before quitting, then we want to
 		// commit at least some of our work (except the last commit,
 		// see above). To do so, we commit every time _before_ we poll.