Skip to content

Commit 6dced71

Browse files
test: update UTs
1 parent 1dd1fae commit 6dced71

3 files changed

Lines changed: 59 additions & 47 deletions

File tree

engine/chunk/chunk_test.go

Lines changed: 49 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -9,34 +9,40 @@ import (
99
"testing"
1010
)
1111

12+
const (
13+
chunkSize = 10
14+
maxPeekSize = 5
15+
smallFileThreshold = int64(20)
16+
)
17+
1218
func TestGetAndPutBuf(t *testing.T) {
13-
chunk := NewChunk()
14-
buf, ok := chunk.GetBuf()
15-
defer chunk.PutBuf(buf)
19+
c := New()
20+
data := []byte("test")
21+
buf, ok := c.GetBuf(data)
22+
defer c.PutBuf(buf)
1623

1724
require.True(t, ok)
18-
require.Equal(t, defaultSize, len(*buf))
25+
require.Equal(t, defaultSize+defaultMaxPeekSize, buf.Cap())
26+
require.Equal(t, string(data), buf.String())
1927
}
2028

21-
func TestGetAndPutPeekBuf(t *testing.T) {
22-
chunk := NewChunk()
23-
data := []byte("test")
24-
buf, ok := chunk.GetPeekBuf(data)
25-
defer chunk.PutPeekBuf(buf)
29+
func TestGetAndPutPeekedBuf(t *testing.T) {
30+
c := New()
31+
window, ok := c.GetPeekedBuf()
32+
defer c.PutPeekedBuf(window)
2633

2734
require.True(t, ok)
28-
require.Equal(t, defaultSize+defaultMaxPeekSize, buf.Cap())
29-
require.Equal(t, string(data), buf.String())
35+
require.Equal(t, defaultSize+defaultMaxPeekSize, len(*window))
3036
}
3137

3238
func TestGetSize(t *testing.T) {
33-
chunk := NewChunk()
34-
require.Equal(t, defaultSize, chunk.GetSize())
39+
c := New()
40+
require.Equal(t, defaultSize, c.GetSize())
3541
}
3642

3743
func TestGetMaxPeekSize(t *testing.T) {
38-
chunk := NewChunk()
39-
require.Equal(t, defaultMaxPeekSize, chunk.GetMaxPeekSize())
44+
c := New()
45+
require.Equal(t, defaultMaxPeekSize, c.GetMaxPeekSize())
4046
}
4147

4248
func TestReadChunk(t *testing.T) {
@@ -66,12 +72,12 @@ func TestReadChunk(t *testing.T) {
6672
{
6773
name: "successful read - peek size exceeded",
6874
reader: strings.NewReader("abc\ndef\nghi\njkl\nmno\npqr\nstu\nvwx\nyz"),
69-
expected: "abc\ndef\nghi\njkl\nmno\npqr\ns",
75+
expected: "abc\ndef\nghi\njkl",
7076
},
7177
{
7278
name: "successful read - multiple lines with consecutives new lines",
73-
reader: strings.NewReader("abc\ndef\n\nghi\n"),
74-
expected: "abc\ndef\n\n",
79+
reader: strings.NewReader("abc\ndef\n\n\n\n\nghi\njkl"),
80+
expected: "abc\ndef\n\n\n",
7581
},
7682
{
7783
name: "multiple lines without consecutives new lines",
@@ -82,11 +88,11 @@ func TestReadChunk(t *testing.T) {
8288

8389
for _, tc := range testCases {
8490
t.Run(tc.name, func(t *testing.T) {
85-
chunk := NewChunkWithSize(5, 20, 20)
86-
reader := bufio.NewReader(tc.reader)
91+
c := New(WithSize(chunkSize), WithMaxPeekSize(maxPeekSize), WithSmallFileThreshold(smallFileThreshold))
92+
reader := bufio.NewReaderSize(tc.reader, chunkSize+maxPeekSize)
8793

8894
// Act
89-
result, err := chunk.ReadChunk(reader, 0)
95+
result, err := c.ReadChunk(reader, 0)
9096
require.ErrorIs(t, err, tc.expectedError)
9197

9298
// Assert
@@ -95,63 +101,62 @@ func TestReadChunk(t *testing.T) {
95101
}
96102
}
97103

98-
func TestReadUntilSafeBoundary(t *testing.T) {
104+
func TestGenerateChunk(t *testing.T) {
99105
// Arrange
100106
testCases := []struct {
101107
name string
102-
reader io.Reader
108+
rawData []byte
103109
expected string
104110
}{
105111
// Current split is fine, exit early.
106112
{
107113
name: "safe original split - LF",
108-
reader: strings.NewReader("abc\n\ndefghijklmnop\n\nqrstuvwxyz"),
109-
expected: "abc\n\n",
114+
rawData: []byte("abc\ndef\n\n\nghijklmnop\n\nqrstuvwxyz"),
115+
expected: "abc\ndef\n\n\n",
110116
},
111117
{
112118
name: "safe original split - CRLF",
113-
reader: strings.NewReader("a\r\n\r\nbcdefghijklmnop\n"),
114-
expected: "a\r\n\r\n",
119+
rawData: []byte("abcdef\r\n\r\nghijklmnop\n"),
120+
expected: "abcdef\r\n\r\n",
115121
},
116-
// Current split is bad, look for a better one.
122+
// Current split is bad, look for a better one
117123
{
118124
name: "safe split - LF",
119-
reader: strings.NewReader("abcdefg\nhijklmnop\n\nqrstuvwxyz"),
120-
expected: "abcdefg\nhijklmnop\n\n",
125+
rawData: []byte("abcdef\nghi\n\njklmnop\n\nqrstuvwxyz"),
126+
expected: "abcdef\nghi\n\n",
121127
},
122128
{
123129
name: "safe split - CRLF",
124-
reader: strings.NewReader("abcdefg\r\nhijklmnop\r\n\r\nqrstuvwxyz"),
125-
expected: "abcdefg\r\nhijklmnop\r\n\r\n",
130+
rawData: []byte("abcdef\r\nghi\r\n\r\njklmnopqrstuvwxyz"),
131+
expected: "abcdef\r\nghi\r\n\r\n",
126132
},
127133
{
128134
name: "safe split - blank line",
129-
reader: strings.NewReader("abcdefg\nhijklmnop\n\t \t\nqrstuvwxyz"),
130-
expected: "abcdefg\nhijklmnop\n\t \t\n",
135+
rawData: []byte("abcdefghi\n\t \t\njklmnopqrstuvwxyz"),
136+
expected: "abcdefghi\n\t \t\n",
131137
},
132-
// Current split is bad, exhaust options.
138+
// Current split is bad, exhaust options
133139
{
134140
name: "no safe split",
135-
reader: strings.NewReader("abcdefg\nhijklmnopqrstuvwxyz"),
136-
expected: "abcdefg\nhijklmnopqrstuvwx",
141+
rawData: []byte("abcdefg\nhijklmnopqrstuvwxyz"),
142+
expected: "abcdefg\nhijklmn",
137143
},
138144
}
139145

140146
for _, tc := range testCases {
141147
t.Run(tc.name, func(t *testing.T) {
142-
buf := make([]byte, 5)
143-
n, err := tc.reader.Read(buf)
148+
c := New(WithSize(chunkSize), WithMaxPeekSize(maxPeekSize), WithSmallFileThreshold(smallFileThreshold))
149+
reader := bufio.NewReaderSize(bytes.NewReader(tc.rawData), c.size+c.maxPeekSize)
150+
peekedBuf := make([]byte, c.size+c.maxPeekSize)
151+
_, err := reader.Read(peekedBuf)
144152
require.NoError(t, err)
145153

146154
// Act
147-
chunk := NewChunkWithSize(5, 20, 20)
148-
reader := bufio.NewReader(tc.reader)
149-
peekBuf := bytes.NewBuffer(buf[:n])
150-
err = chunk.readUntilSafeBoundary(reader, n, peekBuf)
155+
chunkStr, err := c.generateChunk(peekedBuf)
151156
require.NoError(t, err)
152157

153158
// Assert
154-
require.Equal(t, tc.expected, peekBuf.String())
159+
require.Equal(t, tc.expected, chunkStr)
155160
})
156161
}
157162
}

engine/engine.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ func (e *Engine) detectChunks(item plugins.ISourceItem, secretsChannel chan *sec
170170
_ = f.Close()
171171
}()
172172

173-
reader := bufio.NewReaderSize(f, e.chunk.GetMaxPeekSize())
173+
reader := bufio.NewReaderSize(f, e.chunk.GetSize()+e.chunk.GetMaxPeekSize())
174174
totalLines := 0
175175

176176
// Read the file in chunks until EOF

engine/engine_test.go

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -266,8 +266,9 @@ func TestDetectFile(t *testing.T) {
266266
m.chunk.EXPECT().GetFileThreshold().Return(sizeThreshold)
267267
m.chunk.EXPECT().GetSize().Return(chunkSize)
268268
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
269-
m.chunk.EXPECT().GetSize().Return(chunkSize)
270269
m.semaphore.EXPECT().AcquireMemoryWeight(gomock.Any(), chunkWeight).Return(nil)
270+
m.chunk.EXPECT().GetSize().Return(chunkSize)
271+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
271272
m.chunk.EXPECT().ReadChunk(gomock.Any(), gomock.Any()).Return("", assert.AnError)
272273
m.semaphore.EXPECT().ReleaseMemoryWeight(chunkWeight)
273274
},
@@ -283,8 +284,9 @@ func TestDetectFile(t *testing.T) {
283284
m.chunk.EXPECT().GetFileThreshold().Return(sizeThreshold)
284285
m.chunk.EXPECT().GetSize().Return(chunkSize)
285286
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
286-
m.chunk.EXPECT().GetSize().Return(chunkSize)
287287
m.semaphore.EXPECT().AcquireMemoryWeight(gomock.Any(), chunkWeight).Return(nil)
288+
m.chunk.EXPECT().GetSize().Return(chunkSize)
289+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
288290
m.chunk.EXPECT().ReadChunk(gomock.Any(), gomock.Any()).Return("abc\ndef\nghi\njkl\nmno\npqr\nstu\nvw", nil)
289291
m.chunk.EXPECT().ReadChunk(gomock.Any(), gomock.Any()).Return("x\nyz", nil)
290292
m.chunk.EXPECT().ReadChunk(gomock.Any(), gomock.Any()).Return("", io.EOF)
@@ -338,6 +340,7 @@ func TestDetectFile(t *testing.T) {
338340

339341
func TestDetectChunks(t *testing.T) {
340342
chunkSize := 5
343+
maxPeekSize := 20
341344

342345
testCases := []struct {
343346
name string
@@ -351,6 +354,7 @@ func TestDetectChunks(t *testing.T) {
351354
makeFile: func(tmp string) string { return writeTempFile(t, tmp, 0, []byte("password=supersecret\n")) },
352355
mockFunc: func(m *mock) {
353356
m.chunk.EXPECT().GetSize().Return(chunkSize)
357+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
354358
m.chunk.EXPECT().ReadChunk(gomock.Any(), 0).Return("password=supersecret", nil)
355359
m.chunk.EXPECT().ReadChunk(gomock.Any(), 0).Return("", io.EOF)
356360
},
@@ -366,6 +370,7 @@ func TestDetectChunks(t *testing.T) {
366370
makeFile: func(tmp string) string { return writeTempFile(t, tmp, 0, []byte{'P', 'K', 0x03, 0x04}) },
367371
mockFunc: func(m *mock) {
368372
m.chunk.EXPECT().GetSize().Return(chunkSize)
373+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
369374
m.chunk.EXPECT().ReadChunk(gomock.Any(), 0).Return("", fmt.Errorf("skipping file - unsupported file type"))
370375
},
371376
expectedLog: "Skipping file %s: unsupported file type",
@@ -375,6 +380,7 @@ func TestDetectChunks(t *testing.T) {
375380
makeFile: func(tmp string) string { return writeTempFile(t, tmp, 0, []byte("password=supersecret\n")) },
376381
mockFunc: func(m *mock) {
377382
m.chunk.EXPECT().GetSize().Return(chunkSize)
383+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
378384
m.chunk.EXPECT().ReadChunk(gomock.Any(), 0).Return("", io.EOF)
379385
},
380386
},
@@ -383,6 +389,7 @@ func TestDetectChunks(t *testing.T) {
383389
makeFile: func(tmp string) string { return writeTempFile(t, tmp, 0, []byte("password=supersecret\n")) },
384390
mockFunc: func(m *mock) {
385391
m.chunk.EXPECT().GetSize().Return(chunkSize)
392+
m.chunk.EXPECT().GetMaxPeekSize().Return(maxPeekSize)
386393
m.chunk.EXPECT().ReadChunk(gomock.Any(), 0).Return("", assert.AnError)
387394
},
388395
expectedErr: assert.AnError,

0 commit comments

Comments
 (0)