From 9603fe9306592df36c2425aacd226af8e7a0be3d Mon Sep 17 00:00:00 2001
From: Sung Yoon Whang <sungyoon@uber.com>
Date: Thu, 22 Sep 2022 16:14:32 -0700
Subject: [PATCH] Fix decorator not being applied to transient dependencies
 (#941)

There is a bug that only manifests when decorating objects in Fx
where transient dependency types are not properly decorated. For
example, if type A depends on type B, and type B depends on type C
which has a decorator, it fails to recognize that.

This only occurs when all constructors and the decorators for all
these types are provided at the "root" level fx.App.

This happens because fx injects a fake "root" Scope between the
actual root Container (which is where constructors are provided to
by default).

Fixed by making fx stop create this fake root Scope and use the
dig.Container directly.

Fixes #940.
---
 decorate_test.go | 33 +++++++++++++++++++++++++++++++++
 module.go        | 17 +++++++++++++----
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/decorate_test.go b/decorate_test.go
index 49c0a7d75..83c35506f 100644
--- a/decorate_test.go
+++ b/decorate_test.go
@@ -313,6 +313,39 @@ func TestDecorateSuccess(t *testing.T) {
 		)
 		defer app.RequireStart().RequireStop()
 	})
+
+	t.Run("decoration must execute when required by a member of group", func(t *testing.T) {
+		type Drinks interface {
+		}
+		type Coffee struct {
+			Type  string
+			Name  string
+			Price int
+		}
+		type PriceService struct {
+			DefaultPrice int
+		}
+		app := fxtest.New(t,
+			fx.Provide(func() *PriceService {
+				return &PriceService{DefaultPrice: 3}
+			}),
+			fx.Decorate(func(service *PriceService) *PriceService {
+				service.DefaultPrice = 10
+				return service
+			}),
+			fx.Provide(fx.Annotate(func(service *PriceService) Drinks {
+				assert.Equal(t, 10, service.DefaultPrice)
+				return &Coffee{Type: "coffee", Name: "Americano", Price: service.DefaultPrice}
+			}, fx.ResultTags(`group:"drinks"`))),
+			fx.Provide(fx.Annotated{Group: "drinks", Target: func() Drinks {
+				return &Coffee{Type: "coffee", Name: "Cold Brew", Price: 4}
+			}}),
+			fx.Invoke(fx.Annotate(func(drinks []Drinks) {
+				assert.Len(t, drinks, 2)
+			}, fx.ParamTags(`group:"drinks"`))),
+		)
+		defer app.RequireStart().RequireStop()
+	})
 }
 
 func TestDecorateFailure(t *testing.T) {
diff --git a/module.go b/module.go
index c75ba7979..4a769d428 100644
--- a/module.go
+++ b/module.go
@@ -81,7 +81,7 @@ func (o moduleOption) apply(mod *module) {
 type module struct {
 	parent     *module
 	name       string
-	scope      *dig.Scope
+	scope      scope
 	provides   []provide
 	invokes    []invoke
 	decorators []decorator
@@ -89,15 +89,24 @@ type module struct {
 	app        *App
 }
 
+// scope is a private wrapper interface for dig.Container and dig.Scope.
+// We can consider moving this into Fx using type constraints after Go 1.20
+// is released and 1.17 is deprecated.
+type scope interface {
+	Decorate(f interface{}, opts ...dig.DecorateOption) error
+	Invoke(f interface{}, opts ...dig.InvokeOption) error
+	Provide(f interface{}, opts ...dig.ProvideOption) error
+	Scope(name string, opts ...dig.ScopeOption) *dig.Scope
+	String() string
+}
+
 // builds the Scopes using the App's Container. Note that this happens
 // after applyModules' are called because the App's Container needs to
 // be built for any Scopes to be initialized, and applys' should be called
 // before the Container can get initialized.
 func (m *module) build(app *App, root *dig.Container) {
 	if m.parent == nil {
-		m.scope = root.Scope(m.name)
-		// TODO: Once fx.Decorate is in-place,
-		// use the root container instead of subscope.
+		m.scope = root
 	} else {
 		parentScope := m.parent.scope
 		m.scope = parentScope.Scope(m.name)