From e5a78710aa026abfc3f9772ed108ea5e956a634a Mon Sep 17 00:00:00 2001
From: Gongpu Zhu <gongpu@google.com>
Date: Fri, 7 Aug 2020 18:20:00 +0000
Subject: [PATCH] Add callback to kio pipeline execution

---
 kyaml/kio/kio.go      | 12 ++++++++++
 kyaml/kio/kio_test.go | 53 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/kyaml/kio/kio.go b/kyaml/kio/kio.go
index bb999a074..32bdebb32 100644
--- a/kyaml/kio/kio.go
+++ b/kyaml/kio/kio.go
@@ -80,6 +80,15 @@ type Pipeline struct {
 // Execute executes each step in the sequence, returning immediately after encountering
 // any error as part of the Pipeline.
 func (p Pipeline) Execute() error {
+	return p.ExecuteWithCallback(nil)
+}
+
+// PipelineExecuteCallbackFunc defines a callback function that will be called each time a step in the pipeline succeeds.
+type PipelineExecuteCallbackFunc = func(op Filter)
+
+// ExecuteWithCallback executes each step in the sequence, returning immediately after encountering
+// any error as part of the Pipeline. The callback will be called each time a step succeeds.
+func (p Pipeline) ExecuteWithCallback(callback PipelineExecuteCallbackFunc) error {
 	var result []*yaml.RNode
 
 	// read from the inputs
@@ -99,6 +108,9 @@ func (p Pipeline) Execute() error {
 	var err error
 	for i := range p.Filters {
 		op := p.Filters[i]
+		if callback != nil {
+			callback(op)
+		}
 		result, err = op.Filter(result)
 		if len(result) == 0 || err != nil {
 			return errors.Wrap(err)
diff --git a/kyaml/kio/kio_test.go b/kyaml/kio/kio_test.go
index 870cc915a..d51301d9b 100644
--- a/kyaml/kio/kio_test.go
+++ b/kyaml/kio/kio_test.go
@@ -4,9 +4,14 @@
 package kio_test
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"sigs.k8s.io/kustomize/kyaml/yaml"
+
 	. "sigs.k8s.io/kustomize/kyaml/kio"
 )
 
@@ -23,6 +28,50 @@ func TestPipe(t *testing.T) {
 	}
 }
 
-func TestSlice_Write(t *testing.T) {
-
+type mockCallback struct {
+	mock.Mock
+}
+
+func (c *mockCallback) Callback(op Filter) {
+	c.Called(op)
+}
+
+func TestPipelineWithCallback(t *testing.T) {
+	input := ResourceNodeSlice{yaml.MakeNullNode()}
+	noopFilter1 := func(nodes []*yaml.RNode) ([]*yaml.RNode, error) {
+		return nodes, nil
+	}
+	noopFilter2 := func(nodes []*yaml.RNode) ([]*yaml.RNode, error) {
+		return nodes, nil
+	}
+	filters := []Filter{
+		FilterFunc(noopFilter1),
+		FilterFunc(noopFilter2),
+	}
+	p := Pipeline{
+		Inputs:  []Reader{input},
+		Filters: filters,
+		Outputs: []Writer{},
+	}
+
+	callback := mockCallback{}
+	// setup expectations. `Times` means the function is called no more than `times`.
+	callback.On("Callback", mock.Anything).Times(len(filters))
+
+	err := p.ExecuteWithCallback(callback.Callback)
+
+	if !assert.NoError(t, err) {
+		assert.FailNow(t, err.Error())
+	}
+
+	callback.AssertNumberOfCalls(t, "Callback", len(filters))
+
+	// assert filters are called in the order they are defined.
+	for i, filter := range filters {
+		assert.Equal(
+			t,
+			reflect.ValueOf(callback.Calls[i].Arguments[0]).Pointer(),
+			reflect.ValueOf(filter).Pointer(),
+		)
+	}
 }