Fix middleware wrapping

destination_middleware.go

@@ -31,8 +31,6 @@ import (
 	"golang.org/x/time/rate"
 )
 
-var destinationMiddlewareType = reflect.TypeFor[DestinationMiddleware]()
-
 // DestinationMiddleware wraps a Destination and adds functionality to it.
 type DestinationMiddleware interface {
 	Wrap(Destination) Destination
@@ -88,7 +86,21 @@ func DestinationWithMiddleware(d Destination) Destination {
 	if cfgVal.Kind() != reflect.Ptr {
 		panic("The struct returned in Config() must be a pointer")
 	}
-	cfgVal = cfgVal.Elem()
+
+	// Collect all middlewares from the config and wrap the destination with them
+	mw := destinationMiddlewareFromConfigRecursive(cfgVal.Elem())
+
+	// Wrap the middleware in reverse order to preserve the order as specified.
+	for i := len(mw) - 1; i >= 0; i-- {
+		d = mw[i].Wrap(d)
+	}
+
+	return d
+}
+
+func destinationMiddlewareFromConfigRecursive(cfgVal reflect.Value) []DestinationMiddleware {
+	destinationMiddlewareType := reflect.TypeFor[DestinationMiddleware]()
+	cfgType := cfgVal.Type()
 
 	// Collect all middlewares from the config and wrap the destination with them
 	var mw []DestinationMiddleware
@@ -100,19 +112,21 @@ func DestinationWithMiddleware(d Destination) Destination {
 		if field.Kind() != reflect.Ptr {
 			field = field.Addr()
 		}
-		if field.Type().Implements(destinationMiddlewareType) {
+
+		switch {
+		case field.Type().Implements(destinationMiddlewareType):
 			// This is a middleware config, store it.
 			//nolint:forcetypeassert // type checked above with field.Type().Implements()
 			mw = append(mw, field.Interface().(DestinationMiddleware))
-		}
-	}
 
-	// Wrap the middleware in reverse order to preserve the order as specified.
-	for i := len(mw) - 1; i >= 0; i-- {
-		d = mw[i].Wrap(d)
+		case cfgType.Field(i).Anonymous &&
+			cfgType.Field(i).Type.Kind() == reflect.Struct:
+			// This is an embedded struct, dive deeper.
+			mw = append(mw, destinationMiddlewareFromConfigRecursive(field.Elem())...)
+		}
 	}
 
-	return d
+	return mw
 }
 
 // -- DestinationWithBatch -----------------------------------------------------

destination_middleware_test.go

@@ -32,6 +32,28 @@ import (
 	"golang.org/x/time/rate"
 )
 
+func TestDestinationWithMiddleware(t *testing.T) {
+	is := is.New(t)
+
+	ctrl := gomock.NewController(t)
+	src := NewMockDestination(ctrl)
+
+	cfg := struct {
+		DefaultDestinationMiddleware
+	}{}
+	src.EXPECT().Config().Return(&cfg)
+
+	got := DestinationWithMiddleware(src)
+
+	var want Destination = src
+	want = (&DestinationWithSchemaExtraction{}).Wrap(want)
+	want = (&DestinationWithBatch{}).Wrap(want)
+	want = (&DestinationWithRecordFormat{}).Wrap(want)
+	want = (&DestinationWithRateLimit{}).Wrap(want)
+
+	is.Equal(want, got)
+}
+
 // -- DestinationWithBatch -----------------------------------------------------
 
 func TestDestinationWithBatch_Open(t *testing.T) {

destination_test.go

@@ -168,8 +168,6 @@ func TestDestinationPluginAdapter_Run_Write(t *testing.T) {
 }
 
 func TestDestinationPluginAdapter_Run_WriteBatch_Success(t *testing.T) {
-	t.Skip("TODO fix this test")
-
 	is := is.New(t)
 	ctrl := gomock.NewController(t)
 	dst := NewMockDestination(ctrl)
@@ -249,8 +247,6 @@ func TestDestinationPluginAdapter_Run_WriteBatch_Success(t *testing.T) {
 }
 
 func TestDestinationPluginAdapter_Run_WriteBatch_Partial(t *testing.T) {
-	t.Skip("TODO fix this test")
-
 	is := is.New(t)
 	ctrl := gomock.NewController(t)
 	dst := NewMockDestination(ctrl)

source_middleware.go

@@ -31,8 +31,6 @@ import (
 	"github.com/jpillora/backoff"
 )
 
-var sourceMiddlewareType = reflect.TypeFor[SourceMiddleware]()
-
 // SourceMiddleware wraps a Source and adds functionality to it.
 type SourceMiddleware interface {
 	Wrap(Source) Source
@@ -84,7 +82,21 @@ func SourceWithMiddleware(s Source) Source {
 	if cfgVal.Kind() != reflect.Ptr {
 		panic("The struct returned in Config() must be a pointer")
 	}
-	cfgVal = cfgVal.Elem()
+
+	// Collect all middlewares from the config and wrap the source with them
+	mw := sourceMiddlewareFromConfigRecursive(cfgVal.Elem())
+
+	// Wrap the middleware in reverse order to preserve the order as specified.
+	for i := len(mw) - 1; i >= 0; i-- {
+		s = mw[i].Wrap(s)
+	}
+
+	return s
+}
+
+func sourceMiddlewareFromConfigRecursive(cfgVal reflect.Value) []SourceMiddleware {
+	sourceMiddlewareType := reflect.TypeFor[SourceMiddleware]()
+	cfgType := cfgVal.Type()
 
 	// Collect all middlewares from the config and wrap the source with them
 	var mw []SourceMiddleware
@@ -96,19 +108,21 @@ func SourceWithMiddleware(s Source) Source {
 		if field.Kind() != reflect.Ptr {
 			field = field.Addr()
 		}
-		if field.Type().Implements(sourceMiddlewareType) {
+
+		switch {
+		case field.Type().Implements(sourceMiddlewareType):
 			// This is a middleware config, store it.
 			//nolint:forcetypeassert // type checked above with field.Type().Implements()
 			mw = append(mw, field.Interface().(SourceMiddleware))
-		}
-	}
 
-	// Wrap the middleware in reverse order to preserve the order as specified.
-	for i := len(mw) - 1; i >= 0; i-- {
-		s = mw[i].Wrap(s)
+		case cfgType.Field(i).Anonymous &&
+			cfgType.Field(i).Type.Kind() == reflect.Struct:
+			// This is an embedded struct, dive deeper.
+			mw = append(mw, sourceMiddlewareFromConfigRecursive(field.Elem())...)
+		}
 	}
 
-	return s
+	return mw
 }
 
 // -- SourceWithSchemaExtraction -----------------------------------------------

source_middleware_test.go

@@ -31,6 +31,28 @@ import (
 	"go.uber.org/mock/gomock"
 )
 
+func TestSourceWithMiddleware(t *testing.T) {
+	is := is.New(t)
+
+	ctrl := gomock.NewController(t)
+	src := NewMockSource(ctrl)
+
+	cfg := struct {
+		DefaultSourceMiddleware
+	}{}
+	src.EXPECT().Config().Return(&cfg)
+
+	got := SourceWithMiddleware(src)
+
+	var want Source = src
+	want = (&SourceWithSchemaExtraction{}).Wrap(want)
+	want = (&SourceWithSchemaContext{}).Wrap(want)
+	want = (&SourceWithEncoding{}).Wrap(want)
+	want = (&SourceWithBatch{}).Wrap(want)
+
+	is.Equal(want, got)
+}
+
 // -- SourceWithSchemaExtraction -----------------------------------------------
 
 func TestSourceWithSchemaExtraction_SchemaType(t *testing.T) {