7
7
"io"
8
8
"math"
9
9
"strconv"
10
+ "strings"
10
11
"sync"
11
12
"testing"
12
13
"time"
@@ -134,6 +135,10 @@ func TestWriter(t *testing.T) {
134
135
scenario : "writing messages with a small batch byte size" ,
135
136
function : testWriterSmallBatchBytes ,
136
137
},
138
+ {
139
+ scenario : "writing messages with headers" ,
140
+ function : testWriterBatchBytesHeaders ,
141
+ },
137
142
{
138
143
scenario : "setting a non default balancer on the writer" ,
139
144
function : testWriterSetsRightBalancer ,
@@ -449,7 +454,7 @@ func testWriterBatchBytes(t *testing.T) {
449
454
450
455
w := newTestWriter (WriterConfig {
451
456
Topic : topic ,
452
- BatchBytes : 48 ,
457
+ BatchBytes : 50 ,
453
458
BatchTimeout : math .MaxInt32 * time .Second ,
454
459
Balancer : & RoundRobin {},
455
460
})
@@ -458,10 +463,10 @@ func testWriterBatchBytes(t *testing.T) {
458
463
ctx , cancel := context .WithTimeout (context .Background (), 10 * time .Second )
459
464
defer cancel ()
460
465
if err := w .WriteMessages (ctx , []Message {
461
- {Value : []byte ("M0" )}, // 24 Bytes
462
- {Value : []byte ("M1" )}, // 24 Bytes
463
- {Value : []byte ("M2" )}, // 24 Bytes
464
- {Value : []byte ("M3" )}, // 24 Bytes
466
+ {Value : []byte ("M0" )}, // 25 Bytes
467
+ {Value : []byte ("M1" )}, // 25 Bytes
468
+ {Value : []byte ("M2" )}, // 25 Bytes
469
+ {Value : []byte ("M3" )}, // 25 Bytes
465
470
}... ); err != nil {
466
471
t .Error (err )
467
472
return
@@ -592,6 +597,67 @@ func testWriterSmallBatchBytes(t *testing.T) {
592
597
}
593
598
}
594
599
600
+ func testWriterBatchBytesHeaders (t * testing.T ) {
601
+ topic := makeTopic ()
602
+ createTopic (t , topic , 1 )
603
+ defer deleteTopic (t , topic )
604
+
605
+ offset , err := readOffset (topic , 0 )
606
+ if err != nil {
607
+ t .Fatal (err )
608
+ }
609
+
610
+ w := newTestWriter (WriterConfig {
611
+ Topic : topic ,
612
+ BatchBytes : 100 ,
613
+ BatchTimeout : 50 * time .Millisecond ,
614
+ Balancer : & RoundRobin {},
615
+ })
616
+ defer w .Close ()
617
+
618
+ ctx , cancel := context .WithTimeout (context .Background (), 5 * time .Second )
619
+ defer cancel ()
620
+ if err := w .WriteMessages (ctx , []Message {
621
+ {
622
+ Value : []byte ("Hello World 1" ),
623
+ Headers : []Header {
624
+ {Key : "User-Agent" , Value : []byte ("abc/xyz" )},
625
+ },
626
+ },
627
+ {
628
+ Value : []byte ("Hello World 2" ),
629
+ Headers : []Header {
630
+ {Key : "User-Agent" , Value : []byte ("abc/xyz" )},
631
+ },
632
+ },
633
+ }... ); err != nil {
634
+ t .Error (err )
635
+ return
636
+ }
637
+ ws := w .Stats ()
638
+ if ws .Writes != 2 {
639
+ t .Error ("didn't batch messages; Writes: " , ws .Writes )
640
+ return
641
+ }
642
+ msgs , err := readPartition (topic , 0 , offset )
643
+ if err != nil {
644
+ t .Error ("error reading partition" , err )
645
+ return
646
+ }
647
+
648
+ if len (msgs ) != 2 {
649
+ t .Error ("bad messages in partition" , msgs )
650
+ return
651
+ }
652
+
653
+ for _ , m := range msgs {
654
+ if strings .HasPrefix (string (m .Value ), "Hello World" ) {
655
+ continue
656
+ }
657
+ t .Error ("bad messages in partition" , msgs )
658
+ }
659
+ }
660
+
595
661
func testWriterMultipleTopics (t * testing.T ) {
596
662
topic1 := makeTopic ()
597
663
createTopic (t , topic1 , 1 )
0 commit comments