github.com/segmentio/kafka-go@v0.4.48-0.20240318174348-3f6244eb34fd/protocol/record_batch_test.go (about)

     1  package protocol
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"reflect"
     7  	"testing"
     8  	"time"
     9  )
    10  
    11  type memoryRecord struct {
    12  	offset  int64
    13  	time    time.Time
    14  	key     []byte
    15  	value   []byte
    16  	headers []Header
    17  }
    18  
    19  func (m *memoryRecord) Record() Record {
    20  	return Record{
    21  		Offset:  m.offset,
    22  		Time:    m.time,
    23  		Key:     NewBytes(m.key),
    24  		Value:   NewBytes(m.value),
    25  		Headers: m.headers,
    26  	}
    27  }
    28  
    29  func makeRecords(memoryRecords []memoryRecord) []Record {
    30  	records := make([]Record, len(memoryRecords))
    31  	for i, m := range memoryRecords {
    32  		records[i] = m.Record()
    33  	}
    34  	return records
    35  }
    36  
    37  func TestRecordReader(t *testing.T) {
    38  	now := time.Now()
    39  
    40  	records := []memoryRecord{
    41  		{
    42  			offset: 1,
    43  			time:   now,
    44  			key:    []byte("key-1"),
    45  		},
    46  		{
    47  			offset: 2,
    48  			time:   now.Add(time.Millisecond),
    49  			value:  []byte("value-1"),
    50  		},
    51  		{
    52  			offset: 3,
    53  			time:   now.Add(time.Second),
    54  			key:    []byte("key-3"),
    55  			value:  []byte("value-3"),
    56  			headers: []Header{
    57  				{Key: "answer", Value: []byte("42")},
    58  			},
    59  		},
    60  	}
    61  
    62  	r1 := NewRecordReader(makeRecords(records)...)
    63  	r2 := NewRecordReader(makeRecords(records)...)
    64  	assertRecords(t, r1, r2)
    65  }
    66  
    67  func TestMultiRecordReader(t *testing.T) {
    68  	now := time.Now()
    69  
    70  	records := []memoryRecord{
    71  		{
    72  			offset: 1,
    73  			time:   now,
    74  			key:    []byte("key-1"),
    75  		},
    76  		{
    77  			offset: 2,
    78  			time:   now.Add(time.Millisecond),
    79  			value:  []byte("value-1"),
    80  		},
    81  		{
    82  			offset: 3,
    83  			time:   now.Add(time.Second),
    84  			key:    []byte("key-3"),
    85  			value:  []byte("value-3"),
    86  			headers: []Header{
    87  				{Key: "answer", Value: []byte("42")},
    88  			},
    89  		},
    90  	}
    91  
    92  	r1 := NewRecordReader(makeRecords(records)...)
    93  	r2 := MultiRecordReader(
    94  		NewRecordReader(makeRecords(records[:1])...),
    95  		NewRecordReader(makeRecords(records[1:])...),
    96  	)
    97  	assertRecords(t, r1, r2)
    98  }
    99  
   100  func TestControlRecord(t *testing.T) {
   101  	now := time.Now()
   102  
   103  	records := []ControlRecord{
   104  		{
   105  			Offset:  1,
   106  			Time:    now,
   107  			Version: 2,
   108  			Type:    3,
   109  		},
   110  		{
   111  			Offset:  2,
   112  			Time:    now.Add(time.Second),
   113  			Version: 4,
   114  			Type:    5,
   115  			Data:    []byte("Hello World!"),
   116  			Headers: []Header{
   117  				{Key: "answer", Value: []byte("42")},
   118  			},
   119  		},
   120  	}
   121  
   122  	batch := NewControlBatch(records...)
   123  	found := make([]ControlRecord, 0, len(records))
   124  
   125  	for {
   126  		r, err := batch.ReadControlRecord()
   127  		if err != nil {
   128  			if !errors.Is(err, io.EOF) {
   129  				t.Fatal(err)
   130  			}
   131  			break
   132  		}
   133  		found = append(found, *r)
   134  	}
   135  
   136  	if !reflect.DeepEqual(records, found) {
   137  		t.Error("control records mismatch")
   138  	}
   139  }
   140  
   141  func assertRecords(t *testing.T, r1, r2 RecordReader) {
   142  	t.Helper()
   143  
   144  	for {
   145  		rec1, err1 := r1.ReadRecord()
   146  		rec2, err2 := r2.ReadRecord()
   147  
   148  		if err1 != nil || err2 != nil {
   149  			if !errors.Is(err1, err2) {
   150  				t.Error("errors mismatch:")
   151  				t.Log("expected:", err2)
   152  				t.Log("found:   ", err1)
   153  			}
   154  			return
   155  		}
   156  
   157  		if !equalRecords(rec1, rec2) {
   158  			t.Error("records mismatch:")
   159  			t.Logf("expected: %+v", rec2)
   160  			t.Logf("found:    %+v", rec1)
   161  		}
   162  	}
   163  }
   164  
   165  func equalRecords(r1, r2 *Record) bool {
   166  	if r1.Offset != r2.Offset {
   167  		return false
   168  	}
   169  
   170  	if !r1.Time.Equal(r2.Time) {
   171  		return false
   172  	}
   173  
   174  	k1 := readAll(r1.Key)
   175  	k2 := readAll(r2.Key)
   176  
   177  	if !reflect.DeepEqual(k1, k2) {
   178  		return false
   179  	}
   180  
   181  	v1 := readAll(r1.Value)
   182  	v2 := readAll(r2.Value)
   183  
   184  	if !reflect.DeepEqual(v1, v2) {
   185  		return false
   186  	}
   187  
   188  	return reflect.DeepEqual(r1.Headers, r2.Headers)
   189  }
   190  
   191  func readAll(bytes Bytes) []byte {
   192  	if bytes != nil {
   193  		defer bytes.Close()
   194  	}
   195  	b, err := ReadAll(bytes)
   196  	if err != nil {
   197  		panic(err)
   198  	}
   199  	return b
   200  }