Repository: grpc/grpc-go
File: rpc_util_test.go (Go)

Code before:
}

func BenchmarkGZIPCompressor512KiB(b *testing.B) {
	bmCompressor(b, 512*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1MiB(b *testing.B) {
	bmCompressor(b, 1024*1024, NewGZIPCompressor())
}

// mustCompress compresses the input data and returns a BufferSlice.
func mustCompress(t *testing.T, input []byte) mem.BufferSlice {
	t.Helper()
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	if _, err := gz.Write(input); err != nil {
		t.Fatalf("compressInput() failed to write data: %v", err)
	}
	if err := gz.Close(); err != nil {
		t.Fatalf("compressInput() failed to close gzip writer: %v", err)
	}
	compressedData := buf.Bytes()
	return mem.BufferSlice{mem.NewBuffer(&compressedData, nil)}
}

// TestDecompress tests the decompress function with various scenarios, including
// successful decompression, error handling, and edge cases like overflow or
// premature data end. It ensures that the function behaves correctly with different
// inputs, buffer sizes, and error conditions, using the "gzip" compressor for testing.
func (s) TestDecompress(t *testing.T) {
	compressor := encoding.GetCompressor("gzip")

	testCases := []struct {
		name                  string
		input                 mem.BufferSlice
		maxReceiveMessageSize int
		want                  []byte
		wantErr               error
	}{
		{
			name:                  "Decompresses successfully with sufficient buffer size",
			input:                 mustCompress(t, []byte("decompressed data")),
			maxReceiveMessageSize: 50,
			want:                  []byte("decompressed data"),
			wantErr:               nil,
		},
		{
			name:                  "Fails due to exceeding maxReceiveMessageSize",
			input:                 mustCompress(t, []byte("message that is too large")),
			maxReceiveMessageSize: len("message that is too large") - 1,
			want:                  nil,

Review comment (suggestion):
TestDecompress tests the decompress function behaves correctly for following scenarios
- decompress successfully when message is <= maxReceiveMessageSize
- errors when message > maxReceiveMessageSize
- decompress successfully when maxReceiveMessageSize is MaxInt

Code after:
func BenchmarkGZIPCompressor8KiB(b *testing.B) {
	bmCompressor(b, 8*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor64KiB(b *testing.B) {
	bmCompressor(b, 64*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor512KiB(b *testing.B) {
	bmCompressor(b, 512*1024, NewGZIPCompressor())
}

func BenchmarkGZIPCompressor1MiB(b *testing.B) {
	bmCompressor(b, 1024*1024, NewGZIPCompressor())
}

// compressWithDeterministicError compresses the input data and returns a BufferSlice.
func compressWithDeterministicError(t *testing.T, input []byte) mem.BufferSlice {
	t.Helper()
	var buf bytes.Buffer
	gz := gzip.NewWriter(&buf)
	if _, err := gz.Write(input); err != nil {
		t.Fatalf("compressInput() failed to write data: %v", err)
	}
	if err := gz.Close(); err != nil {
		t.Fatalf("compressInput() failed to close gzip writer: %v", err)
	}
	compressedData := buf.Bytes()
	return mem.BufferSlice{mem.NewBuffer(&compressedData, nil)}
}

// MockDecompressor is a mock implementation of a decompressor used for testing purposes.
// It simulates decompression behavior, returning either decompressed data or an error based on the ShouldError flag.
type MockDecompressor struct {
	ShouldError bool // Flag to control whether the decompression should simulate an error.
}

// Do simulates decompression. It returns a predefined error if ShouldError is true,
// or a fixed set of decompressed data if ShouldError is false.
func (m *MockDecompressor) Do(_ io.Reader) ([]byte, error) {
	if m.ShouldError {
		return nil, errors.New(decompressionErrorMsg)
	}
	return []byte(defaultDecompressedData), nil
}

// Type returns the string identifier for the MockDecompressor.
func (m *MockDecompressor) Type() string {
	return "MockDecompressor"
}