kgo: add test around purging then adding direct partitions

Author: twmb

Diff for: pkg/kgo/consumer_direct_test.go

@@ -45,7 +45,7 @@ func TestIssue337(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
 	defer cancel()
 	for {
 		fs := cl.PollFetches(ctx)
@@ -60,3 +60,55 @@ func TestIssue337(t *testing.T) {
 		}
 	}
 }
+
+func TestDirectPartitionPurge(t *testing.T) {
+	topic, cleanup := tmpTopicPartitions(t, 2)
+	defer cleanup()
+
+	cl, _ := NewClient(
+		DefaultProduceTopic(topic),
+		RecordPartitioner(ManualPartitioner()),
+		ConsumePartitions(map[string]map[int32]Offset{
+			topic: {0: NewOffset().At(0)},
+		}),
+	)
+
+	if err := cl.ProduceSync(context.Background(),
+		&Record{Partition: 0, Value: []byte("foo")},
+		&Record{Partition: 1, Value: []byte("bar")},
+	).FirstErr(); err != nil {
+		t.Fatal(err)
+	}
+	cl.PurgeTopicsFromClient(topic)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	fs := cl.PollFetches(ctx)
+	cancel()
+	if err := fs.Err0(); err != context.DeadlineExceeded {
+		t.Fatal("unexpected success when expecting context.DeadlineExceeded")
+	}
+
+	cl.AddConsumeTopics(topic)
+	ctx, cancel = context.WithTimeout(context.Background(), 3*time.Second)
+
+	exp := map[string]bool{
+		"foo": true,
+		"bar": true,
+	}
+	for {
+		fs := cl.PollFetches(ctx)
+		if err := fs.Err0(); err == context.DeadlineExceeded {
+			break
+		}
+		fs.EachRecord(func(r *Record) {
+			v := string(r.Value)
+			if !exp[v] {
+				t.Errorf("saw unexpected value %v", v)
+			}
+			delete(exp, v)
+		})
+	}
+	if len(exp) > 0 {
+		t.Errorf("did not see expected values %v", exp)
+	}
+}