update to remove totalsize.

trufflesecurity · Aug 2, 2023 · d185482 · d185482
1 parent fbce149
commit d185482
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 22 deletions.
diff --git a/pkg/sources/chunker.go b/pkg/sources/chunker.go
@@ -65,14 +65,6 @@ func WithChunkSize(size int) ConfigOption {
 	}
 }
 
-// WithTotalChunkSize sets the total chunk size.
-// This is the chunk size plus the peek size.
-func WithTotalChunkSize(size int) ConfigOption {
-	return func(c *chunkReaderConfig) {
-		c.totalSize = size
-	}
-}
-
 // WithPeekSize sets the peek size.
 func WithPeekSize(size int) ConfigOption {
 	return func(c *chunkReaderConfig) {
@@ -97,13 +89,14 @@ func applyOptions(opts []ConfigOption) *chunkReaderConfig {
 	config := &chunkReaderConfig{
 		chunkSize: ChunkSize,      // default
 		totalSize: TotalChunkSize, // default
-		peekSize:  PeekSize,       // default
 	}
 
 	for _, opt := range opts {
 		opt(config)
 	}
 
+	config.totalSize = config.chunkSize + config.peekSize
+
 	return config
 }
 

diff --git a/pkg/sources/chunker_test.go b/pkg/sources/chunker_test.go
@@ -73,7 +73,6 @@ func TestNewChunkedReader(t *testing.T) {
 		input      string
 		chunkSize  int
 		peekSize   int
-		totalSize  int
 		wantChunks []string
 		wantErr    bool
 	}{
@@ -82,7 +81,6 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      "example input",
 			chunkSize:  ChunkSize,
 			peekSize:   PeekSize,
-			totalSize:  TotalChunkSize,
 			wantChunks: []string{"example input"},
 			wantErr:    false,
 		},
@@ -91,7 +89,6 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      "",
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
 			wantChunks: []string{},
 			wantErr:    false,
 		},
@@ -100,7 +97,6 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      "small data",
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
 			wantChunks: []string{"small data"},
 			wantErr:    false,
 		},
@@ -109,7 +105,6 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      strings.Repeat("a", 1024),
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
 			wantChunks: []string{strings.Repeat("a", 1024)},
 			wantErr:    false,
 		},
@@ -118,25 +113,22 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      strings.Repeat("a", 1536),
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
-			wantChunks: []string{strings.Repeat("a", 1024), strings.Repeat("a", 512)},
+			wantChunks: []string{strings.Repeat("a", 1536), strings.Repeat("a", 512)},
 			wantErr:    false,
 		},
 		{
 			name:       "EOF during peeking",
 			input:      strings.Repeat("a", 1300),
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
-			wantChunks: []string{strings.Repeat("a", 1024), strings.Repeat("a", 276)},
+			wantChunks: []string{strings.Repeat("a", 1300), strings.Repeat("a", 276)},
 			wantErr:    false,
 		},
 		{
 			name:       "EOF during reading",
 			input:      strings.Repeat("a", 512),
 			chunkSize:  1024,
 			peekSize:   512,
-			totalSize:  1024 + 512,
 			wantChunks: []string{strings.Repeat("a", 512)},
 			wantErr:    false,
 		},
@@ -145,7 +137,6 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      strings.Repeat("a", 2048),
 			chunkSize:  1024,
 			peekSize:   1024,
-			totalSize:  1024 + 1024,
 			wantChunks: []string{strings.Repeat("a", 2048), strings.Repeat("a", 1024)},
 			wantErr:    false,
 		},
@@ -154,15 +145,14 @@ func TestNewChunkedReader(t *testing.T) {
 			input:      strings.Repeat("a", 4096),
 			chunkSize:  1024,
 			peekSize:   1024,
-			totalSize:  1024 + 1024,
 			wantChunks: []string{strings.Repeat("a", 2048), strings.Repeat("a", 2048), strings.Repeat("a", 2048), strings.Repeat("a", 1024)},
 			wantErr:    false,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			readerFunc := NewChunkReader(WithChunkSize(tt.chunkSize), WithPeekSize(tt.peekSize), WithTotalChunkSize(tt.totalSize))
+			readerFunc := NewChunkReader(WithChunkSize(tt.chunkSize), WithPeekSize(tt.peekSize))
 			reader := strings.NewReader(tt.input)
 			ctx := context.Background()
 			dataChan, errChan := readerFunc(ctx, reader)