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

Code before:

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

// compressInput compresses the input data and returns a BufferSlice.
func compressInput(t *testing.T, input []byte) mem.BufferSlice {
	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 TestDecompress(t *testing.T) {
	c := encoding.GetCompressor("gzip")

	tests := []struct {
		name                  string
		compressor            encoding.Compressor
		input                 []byte
		maxReceiveMessageSize int
		want                  []byte
		wantErr               error
	}{
		{
			name:                  "Decompresses successfully with sufficient buffer size",
			compressor:            c,
			input:                 []byte("decompressed data"),
			maxReceiveMessageSize: 50,
			want:                  []byte("decompressed data"),
			wantErr:               nil,
		},
		{
			name:                  "Fails due to exceeding maxReceiveMessageSize",
			compressor:            c,
			input:                 []byte("small message that is too large"),
			maxReceiveMessageSize: 5,
			want:                  nil,
			wantErr:               errMaxMessageSizeExceeded,

Review comment (suggestion):
Please define this as a method on the tester which does extra things like goroutine leak detection.

```suggestion
func (s) TestDecompress(t *testing.T) {
```

Code after:
	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"
}

// 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
// errors when the decompressed message has an invalid format