/
testutils.go
156 lines (139 loc) · 3.37 KB
/
testutils.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
package mcap
import (
"bytes"
"encoding/binary"
"hash/crc32"
"io"
"testing"
"github.com/klauspost/compress/zstd"
"github.com/pierrec/lz4/v4"
"github.com/stretchr/testify/assert"
)
func encodedUint16(x uint16) []byte {
buf := make([]byte, 2)
binary.LittleEndian.PutUint16(buf, x)
return buf
}
func encodedUint32(x uint32) []byte {
buf := make([]byte, 4)
binary.LittleEndian.PutUint32(buf, x)
return buf
}
func encodedUint64(x uint64) []byte {
buf := make([]byte, 8)
binary.LittleEndian.PutUint64(buf, x)
return buf
}
func prefixedString(s string) []byte {
buf := make([]byte, len(s)+4)
binary.LittleEndian.PutUint32(buf, uint32(len(s)))
copy(buf[4:], s)
return buf
}
func prefixedBytes(s []byte) []byte {
buf := make([]byte, len(s)+4)
binary.LittleEndian.PutUint32(buf, uint32(len(s)))
copy(buf[4:], s)
return buf
}
func flatten(slices ...[]byte) []byte {
var flattened []byte
for _, s := range slices {
flattened = append(flattened, s...)
}
return flattened
}
func file(records ...[]byte) []byte {
var file [][]byte
file = append(file, Magic)
file = append(file, records...)
file = append(file, Magic)
return flatten(file...)
}
func footer() []byte {
buf := make([]byte, 9)
buf[0] = byte(OpFooter)
return buf
}
func header() []byte {
buf := make([]byte, 9)
buf[0] = byte(OpHeader)
return buf
}
func channelInfo() []byte {
buf := make([]byte, 9)
buf[0] = byte(OpChannel)
return buf
}
func message() []byte {
buf := make([]byte, 9)
buf[0] = byte(OpMessage)
return buf
}
func chunk(t *testing.T, compression CompressionFormat, includeCRC bool, records ...[]byte) []byte {
data := flatten(records...)
buf := &bytes.Buffer{}
switch compression {
case CompressionZSTD:
w, err := zstd.NewWriter(buf)
if err != nil {
t.Errorf("failed to create zstd writer: %s", err)
}
_, err = io.Copy(w, bytes.NewReader(data))
assert.Nil(t, err)
w.Close()
case CompressionLZ4:
w := lz4.NewWriter(buf)
_, err := io.Copy(w, bytes.NewReader(data))
assert.Nil(t, err)
w.Close()
case CompressionNone:
_, err := buf.Write(data)
assert.Nil(t, err)
default:
_, err := buf.Write(data) // unrecognized compression
assert.Nil(t, err)
}
compressionLen := len(compression)
compressedLen := buf.Len()
uncompressedLen := len(data)
msglen := uint64(8 + 8 + 8 + 4 + 4 + compressionLen + 8 + compressedLen)
record := make([]byte, msglen+9)
offset, err := putByte(record, byte(OpChunk))
assert.Nil(t, err)
offset += putUint64(record[offset:], msglen)
offset += putUint64(record[offset:], 0) // start
offset += putUint64(record[offset:], 1e9) // end
offset += putUint64(record[offset:], uint64(uncompressedLen))
var crc uint32
if includeCRC {
sum := crc32.NewIEEE()
_, _ = sum.Write(data)
crc = sum.Sum32()
} else {
crc = 0
}
offset += putUint32(record[offset:], crc)
offset += putPrefixedString(record[offset:], string(compression))
offset += putUint64(record[offset:], uint64(buf.Len()))
_ = copy(record[offset:], buf.Bytes())
return record
}
func record(op OpCode) []byte {
buf := make([]byte, 9)
buf[0] = byte(op)
return buf
}
func attachment() []byte {
recordLen := 9 + // opcode + record length
8 + // record time
8 + // create time
4 + // attachment name length
4 + // media type length
8 + // data size
4 // crc length
buf := make([]byte, recordLen)
buf[0] = byte(OpAttachment)
putUint64(buf[1:], uint64(recordLen-9))
return buf
}