Check number of attempts in test

huangfu · Jun 15, 2020 · 2695e01 · 2695e01
1 parent 55e867e
commit 2695e01
Showing 1 changed file with 13 additions and 5 deletions.
diff --git a/writer_test.go b/writer_test.go
@@ -140,14 +140,17 @@ func TestValidateWriter(t *testing.T) {
 	}
 }
 
-type fakeWriter struct{}
+type fakeWriter struct {
+	attempts int
+}
 
 func (f *fakeWriter) messages() chan<- writerMessage {
 	ch := make(chan writerMessage, 1)
 
 	go func() {
 		for {
 			msg := <-ch
+			f.attempts++
 			msg.res <- &writerError{
 				err: errors.New("bad attempt"),
 			}
@@ -157,20 +160,21 @@ func (f *fakeWriter) messages() chan<- writerMessage {
 	return ch
 }
 
-func (f *fakeWriter) close() {
-
-}
+func (f *fakeWriter) close() {}
 
 func testWriterMaxAttemptsErr(t *testing.T) {
 	const topic = "test-writer-2"
+	const maxAttempts = 3
+
+	var fw fakeWriter
 
 	createTopic(t, topic, 1)
 	w := newTestWriter(WriterConfig{
 		Topic:       topic,
 		MaxAttempts: 1,
 		Balancer:    &RoundRobin{},
 		newPartitionWriter: func(p int, config WriterConfig, stats *writerStats) partitionWriter {
-			return &fakeWriter{}
+			return &fw
 		},
 	})
 	defer w.Close()
@@ -186,6 +190,10 @@ func testWriterMaxAttemptsErr(t *testing.T) {
 			return
 		}
 	}
+
+	if fw.attempts != maxAttempts {
+		t.Errorf("got %d attempts, want %d", fw.attempts, maxAttempts)
+	}
 }
 
 func testWriterMaxBytes(t *testing.T) {