Fix WithCompression to behave like WithAcceptCompression (#493)

Update doc comments and add tests, too
connectrpc · Apr 11, 2023 · c241c0e · c241c0e
1 parent c80677b
commit c241c0e
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 19 deletions.
diff --git a/compression.go b/compression.go
@@ -65,6 +65,9 @@ func newCompressionPool(
 	newDecompressor func() Decompressor,
 	newCompressor func() Compressor,
 ) *compressionPool {
+	if newDecompressor == nil && newCompressor == nil {
+		return nil
+	}
 	return &compressionPool{
 		decompressors: sync.Pool{
 			New: func() any { return newDecompressor() },

diff --git a/compression_test.go b/compression_test.go
@@ -54,3 +54,110 @@ func TestAcceptEncodingOrdering(t *testing.T) {
 	_, _ = client.CallUnary(context.Background(), NewRequest(&emptypb.Empty{}))
 	assert.True(t, called)
 }
+
+func TestClientCompressionOptionTest(t *testing.T) {
+	t.Parallel()
+	const testURL = "http://foo.bar.com/service/method"
+
+	checkPools := func(t *testing.T, config *clientConfig) {
+		t.Helper()
+		assert.Equal(t, len(config.CompressionNames), len(config.CompressionPools))
+		for _, name := range config.CompressionNames {
+			pool := config.CompressionPools[name]
+			assert.NotNil(t, pool)
+		}
+	}
+	dummyDecompressCtor := func() Decompressor { return nil }
+	dummyCompressCtor := func() Compressor { return nil }
+
+	t.Run("defaults", func(t *testing.T) {
+		t.Parallel()
+		config, err := newClientConfig(testURL, nil)
+		assert.Nil(t, err)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithAcceptCompression", func(t *testing.T) {
+		t.Parallel()
+		opts := []ClientOption{WithAcceptCompression("foo", dummyDecompressCtor, dummyCompressCtor)}
+		config, err := newClientConfig(testURL, opts)
+		assert.Nil(t, err)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip, "foo"})
+		checkPools(t, config)
+	})
+	t.Run("WithAcceptCompression-empty-name-noop", func(t *testing.T) {
+		t.Parallel()
+		opts := []ClientOption{WithAcceptCompression("", dummyDecompressCtor, dummyCompressCtor)}
+		config, err := newClientConfig(testURL, opts)
+		assert.Nil(t, err)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithAcceptCompression-nil-ctors-noop", func(t *testing.T) {
+		t.Parallel()
+		opts := []ClientOption{WithAcceptCompression("foo", nil, nil)}
+		config, err := newClientConfig(testURL, opts)
+		assert.Nil(t, err)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithAcceptCompression-nil-ctors-unregisters", func(t *testing.T) {
+		t.Parallel()
+		opts := []ClientOption{WithAcceptCompression("gzip", nil, nil)}
+		config, err := newClientConfig(testURL, opts)
+		assert.Nil(t, err)
+		assert.Equal(t, config.CompressionNames, nil)
+		checkPools(t, config)
+	})
+}
+
+func TestHandlerCompressionOptionTest(t *testing.T) {
+	t.Parallel()
+	const testProc = "/service/method"
+
+	checkPools := func(t *testing.T, config *handlerConfig) {
+		t.Helper()
+		assert.Equal(t, len(config.CompressionNames), len(config.CompressionPools))
+		for _, name := range config.CompressionNames {
+			pool := config.CompressionPools[name]
+			assert.NotNil(t, pool)
+		}
+	}
+	dummyDecompressCtor := func() Decompressor { return nil }
+	dummyCompressCtor := func() Compressor { return nil }
+
+	t.Run("defaults", func(t *testing.T) {
+		t.Parallel()
+		config := newHandlerConfig(testProc, nil)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithCompression", func(t *testing.T) {
+		t.Parallel()
+		opts := []HandlerOption{WithCompression("foo", dummyDecompressCtor, dummyCompressCtor)}
+		config := newHandlerConfig(testProc, opts)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip, "foo"})
+		checkPools(t, config)
+	})
+	t.Run("WithCompression-empty-name-noop", func(t *testing.T) {
+		t.Parallel()
+		opts := []HandlerOption{WithCompression("", dummyDecompressCtor, dummyCompressCtor)}
+		config := newHandlerConfig(testProc, opts)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithCompression-nil-ctors-noop", func(t *testing.T) {
+		t.Parallel()
+		opts := []HandlerOption{WithCompression("foo", nil, nil)}
+		config := newHandlerConfig(testProc, opts)
+		assert.Equal(t, config.CompressionNames, []string{compressionGzip})
+		checkPools(t, config)
+	})
+	t.Run("WithCompression-nil-ctors-unregisters", func(t *testing.T) {
+		t.Parallel()
+		opts := []HandlerOption{WithCompression("gzip", nil, nil)}
+		config := newHandlerConfig(testProc, opts)
+		assert.Equal(t, config.CompressionNames, nil)
+		checkPools(t, config)
+	})
+}
diff --git a/option.go b/option.go
@@ -43,14 +43,13 @@ type ClientOption interface {
 // Clients accept gzipped responses by default, using a compressor backed by the
 // standard library's [gzip] package with the default compression level. Use
 // [WithSendGzip] to compress requests with gzip.
+//
+// Calling WithAcceptCompression with an empty name is a no-op.
 func WithAcceptCompression(
 	name string,
 	newDecompressor func() Decompressor,
 	newCompressor func() Compressor,
 ) ClientOption {
-	if newDecompressor == nil && newCompressor == nil {
-		return &compressionOption{Name: name}
-	}
 	return &compressionOption{
 		Name:            name,
 		CompressionPool: newCompressionPool(newDecompressor, newCompressor),
@@ -93,7 +92,7 @@ func WithSendCompression(name string) ClientOption {
 
 // WithSendGzip configures the client to gzip requests. Since clients have
 // access to a gzip compressor by default, WithSendGzip doesn't require
-// [WithSendCompresion].
+// [WithSendCompression].
 //
 // Some servers don't support gzip, so clients default to sending uncompressed
 // requests.
@@ -116,9 +115,11 @@ type HandlerOption interface {
 // compressors and decompressors.
 //
 // By default, handlers support gzip using the standard library's
-// [compress/gzip] package at the default compression level.
+// [compress/gzip] package at the default compression level. To remove support for
+// a previously-registered compression algorithm, use WithCompression with nil
+// decompressor and compressor constructors.
 //
-// Calling WithCompression with an empty name or nil constructors is a no-op.
+// Calling WithCompression with an empty name is a no-op.
 func WithCompression(
 	name string,
 	newDecompressor func() Decompressor,
@@ -350,31 +351,31 @@ type compressionOption struct {
 }
 
 func (o *compressionOption) applyToClient(config *clientConfig) {
+	o.apply(&config.CompressionNames, config.CompressionPools)
+}
+
+func (o *compressionOption) applyToHandler(config *handlerConfig) {
+	o.apply(&config.CompressionNames, config.CompressionPools)
+}
+
+func (o *compressionOption) apply(configuredNames *[]string, configuredPools map[string]*compressionPool) {
 	if o.Name == "" {
 		return
 	}
 	if o.CompressionPool == nil {
-		delete(config.CompressionPools, o.Name)
+		delete(configuredPools, o.Name)
 		var names []string
-		for _, name := range config.CompressionNames {
+		for _, name := range *configuredNames {
 			if name == o.Name {
 				continue
 			}
 			names = append(names, name)
 		}
-		config.CompressionNames = names
-		return
-	}
-	config.CompressionPools[o.Name] = o.CompressionPool
-	config.CompressionNames = append(config.CompressionNames, o.Name)
-}
-
-func (o *compressionOption) applyToHandler(config *handlerConfig) {
-	if o.Name == "" || o.CompressionPool == nil {
+		*configuredNames = names
 		return
 	}
-	config.CompressionPools[o.Name] = o.CompressionPool
-	config.CompressionNames = append(config.CompressionNames, o.Name)
+	configuredPools[o.Name] = o.CompressionPool
+	*configuredNames = append(*configuredNames, o.Name)
 }
 
 type compressMinBytesOption struct {