From 9f88fb65560e76231e8d07db854cb91a85828abd Mon Sep 17 00:00:00 2001
From: Steven Borrelli <steve@borrelli.org>
Date: Thu, 3 Oct 2024 12:24:27 +0200
Subject: [PATCH] move merge logic out of main function

Signed-off-by: Steven Borrelli <steve@borrelli.org>
---
 context.go | 30 ++++++++++++++++++++++++++++++
 fn.go      | 27 ++++++++-------------------
 2 files changed, 38 insertions(+), 19 deletions(-)
 create mode 100644 context.go

diff --git a/context.go b/context.go
new file mode 100644
index 0000000..d8b86ba
--- /dev/null
+++ b/context.go
@@ -0,0 +1,30 @@
+package main
+
+import (
+	"dario.cat/mergo"
+	"github.com/crossplane/function-sdk-go/errors"
+	fnv1beta1 "github.com/crossplane/function-sdk-go/proto/v1beta1"
+	"github.com/crossplane/function-sdk-go/request"
+	"github.com/crossplane/function-sdk-go/resource"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
+)
+
+type Context map[string]any
+
+// MergeContextKey merges existing Context at a key with context data val
+func (f *Function) MergeContextKey(key string, val map[string]interface{}, req *fnv1beta1.RunFunctionRequest) (*unstructured.Unstructured, error) {
+	// Check if key is already defined in the context and merge fields
+	var mergedContext *unstructured.Unstructured
+	if v, ok := request.GetContextKey(req, key); ok {
+		mergedContext = &unstructured.Unstructured{}
+		if err := resource.AsObject(v.GetStructValue(), mergedContext); err != nil {
+			return mergedContext, errors.Wrapf(err, "cannot get Composition environment from %T context key %q", req, key)
+		}
+		f.log.Debug("Loaded Existing Function Context", "context-key", key)
+		if err := mergo.Merge(&mergedContext.Object, val); err != nil {
+			return mergedContext, errors.Wrapf(err, "cannot merge data %T at context key %q", req, key)
+		}
+		return mergedContext, nil
+	}
+	return &unstructured.Unstructured{Object: val}, nil
+}
diff --git a/fn.go b/fn.go
index cbaa7c6..6df38f6 100644
--- a/fn.go
+++ b/fn.go
@@ -194,7 +194,7 @@ func (f *Function) RunFunction(_ context.Context, req *fnv1beta1.RunFunctionRequ
 			case "Context":
 				contextData := make(map[string]any)
 				if err = cd.Resource.GetValueInto("data", &contextData); err != nil {
-					response.Fatal(rsp, errors.Wrap(err, "cannot get contexts"))
+					response.Fatal(rsp, errors.Wrap(err, "cannot get Contexts from input"))
 					return rsp, nil
 				}
 				for key, data := range contextData {
@@ -203,27 +203,16 @@ func (f *Function) RunFunction(_ context.Context, req *fnv1beta1.RunFunctionRequ
 						response.Fatal(rsp, errors.Wrapf(err, "cannot convert Context from %T context key %q", req, key))
 						return rsp, nil
 					}
-					// Check if key is already defined in the context
-					var inputEnv *unstructured.Unstructured
-					if v, ok := request.GetContextKey(req, key); ok {
-						inputEnv = &unstructured.Unstructured{}
-						if err := resource.AsObject(v.GetStructValue(), inputEnv); err != nil {
-							response.Fatal(rsp, errors.Wrapf(err, "cannot get Composition environment from %T context key %q", req, key))
-							return rsp, nil
-						}
-						f.log.Debug("Loaded Existing Composition environment from Function context", "context-key", key)
-						if err := mergo.Merge(&inputEnv.Object, val); err != nil {
-							response.Fatal(rsp, errors.Wrapf(err, "cannot merge data %T at context key %q", req, key))
-							return rsp, nil
-						}
-					} else {
-						inputEnv = &unstructured.Unstructured{Object: val}
+					mergedCtx, err := f.MergeContextKey(key, val, req)
+					if err != nil {
+						response.Fatal(rsp, errors.Wrapf(err, "cannot merge Context"))
+						return rsp, nil
 					}
 
-					if inputEnv.GroupVersionKind().Empty() {
-						inputEnv.SetGroupVersionKind(schema.GroupVersionKind{Group: "internal.crossplane.io", Kind: "Environment", Version: "v1alpha1"})
+					if mergedCtx.GroupVersionKind().Empty() {
+						mergedCtx.SetGroupVersionKind(schema.GroupVersionKind{Group: "internal.crossplane.io", Kind: "Environment", Version: "v1alpha1"})
 					}
-					v, err := resource.AsStruct(inputEnv)
+					v, err := resource.AsStruct(mergedCtx)
 					if err != nil {
 						response.Fatal(rsp, errors.Wrap(err, "cannot convert Context to protobuf Struct well-known type"))
 						return rsp, nil