From 136a245299fad97615d6c28eec46ba4502947140 Mon Sep 17 00:00:00 2001
From: Pavel Smejkal <pavel.smejkal@outreach.io>
Date: Wed, 26 Jun 2024 12:20:23 +0200
Subject: [PATCH] feat: wrap, definers

---
 container.go |  5 ++++-
 plumber.go   | 12 ++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/container.go b/container.go
index dfb4ea1..38f6341 100644
--- a/container.go
+++ b/container.go
@@ -39,9 +39,12 @@ func (c *Container) Close() error {
 }
 
 // DefineContainers defines supplied containers using given config and root container
-func DefineContainers[C, CF any](ctx context.Context, cfg CF, root C, containers ...interface {
+func DefineContainers[C, CF any](ctx context.Context, cfg CF, definers []func(context.Context, CF, C), root C, containers ...interface {
 	Define(context.Context, CF, C)
 }) C {
+	for _, d := range definers {
+		d(ctx, cfg, root)
+	}
 	for _, d := range containers {
 		d.Define(ctx, cfg, root)
 	}
diff --git a/plumber.go b/plumber.go
index 98e9029..0aafc9d 100644
--- a/plumber.go
+++ b/plumber.go
@@ -73,6 +73,7 @@ type D[T any] struct {
 	resolve   func()
 	deps      []Dependency
 	listeners []func()
+	wrappers  []func(T) T
 }
 
 // String return names of underlaying type
@@ -89,6 +90,9 @@ func (d *D[T]) define(resolve func()) {
 			resolve()
 			d.resolved = true
 			d.resolving = false
+			for _, w := range d.wrappers {
+				d.value = w(d.value)
+			}
 			for _, l := range d.listeners {
 				l()
 			}
@@ -193,6 +197,14 @@ func (d *D[T]) WhenResolved(callback func()) *D[T] {
 	return d
 }
 
+// Wrap registers a wrapping callback that will be triggered when dependency is resolved
+// The callback allows to augment the original value. Wrapping should be used mostly to
+// redefine the dependency for a different test environments
+func (d *D[T]) Wrap(wrappers ...func(T) T) *D[T] {
+	d.wrappers = append(d.wrappers, wrappers...)
+	return d
+}
+
 // R represents a runnable dependency wrapper
 // It is meant to be supplied into the Pipeline()
 type R[T any] struct {