From c3db3c3b52e50cab618272361129fd277bd993ec Mon Sep 17 00:00:00 2001
From: "Maximilian Blatt (external expert on behalf of DB Netz)"
 <maximilian.blatt-extern@deutschebahn.com>
Date: Thu, 14 Dec 2023 12:20:56 +0100
Subject: [PATCH] fix(eks): Update NodeGroup taints

Signed-off-by: Maximilian Blatt (external expert on behalf of DB Netz) <maximilian.blatt-extern@deutschebahn.com>
---
 pkg/clients/eks/nodegroup.go      | 51 +++++++++++++++++++++-
 pkg/clients/eks/nodegroup_test.go | 72 +++++++++++++++++++++++++++++++
 2 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/pkg/clients/eks/nodegroup.go b/pkg/clients/eks/nodegroup.go
index 3fe9dca566..45eccf1116 100644
--- a/pkg/clients/eks/nodegroup.go
+++ b/pkg/clients/eks/nodegroup.go
@@ -181,10 +181,55 @@ func GenerateUpdateNodeGroupConfigInput(name string, p *manualv1alpha1.NodeGroup
 			MaxUnavailablePercentage: p.UpdateConfig.MaxUnavailablePercentage,
 		}
 	}
-	// TODO(muvaf): Add support for updating taints.
+
+	if p.Taints != nil {
+		u.Taints = generateUpdateTaintsPayload(p.Taints, ng.Taints)
+	}
 	return u
 }
 
+func generateUpdateTaintsPayload(spec []manualv1alpha1.Taint, current []ekstypes.Taint) *ekstypes.UpdateTaintsPayload {
+	res := &ekstypes.UpdateTaintsPayload{}
+
+	curMap := map[string]manualv1alpha1.Taint{}
+	for _, t := range current {
+		if t.Key == nil {
+			continue
+		}
+		curMap[*t.Key] = manualv1alpha1.Taint{
+			Effect: string(t.Effect),
+			Key:    t.Key,
+			Value:  t.Value,
+		}
+	}
+
+	specMap := map[string]any{}
+	for _, st := range spec {
+		if st.Key == nil {
+			continue
+		}
+		specMap[*st.Key] = nil
+
+		ct, exists := curMap[*st.Key]
+		if !exists || !cmp.Equal(st, ct) {
+			res.AddOrUpdateTaints = append(res.AddOrUpdateTaints, ekstypes.Taint{
+				Effect: ekstypes.TaintEffect(st.Effect),
+				Key:    st.Key,
+				Value:  st.Value,
+			})
+		}
+	}
+	for _, ct := range current {
+		if ct.Key == nil {
+			continue
+		}
+		if _, exists := specMap[*ct.Key]; !exists {
+			res.RemoveTaints = append(res.RemoveTaints, ct)
+		}
+	}
+	return res
+}
+
 // GenerateNodeGroupObservation is used to produce manualv1alpha1.NodeGroupObservation
 // from eks.Nodegroup.
 func GenerateNodeGroupObservation(ng *ekstypes.Nodegroup) manualv1alpha1.NodeGroupObservation { //nolint:gocyclo
@@ -351,6 +396,10 @@ func IsNodeGroupUpToDate(p *manualv1alpha1.NodeGroupParameters, ng *ekstypes.Nod
 		}
 		return true
 	}
+	taints := generateUpdateTaintsPayload(p.Taints, ng.Taints)
+	if len(taints.AddOrUpdateTaints) > 0 || len(taints.RemoveTaints) > 0 {
+		return false
+	}
 	return false
 }
 
diff --git a/pkg/clients/eks/nodegroup_test.go b/pkg/clients/eks/nodegroup_test.go
index 2780c23965..be76b4e760 100644
--- a/pkg/clients/eks/nodegroup_test.go
+++ b/pkg/clients/eks/nodegroup_test.go
@@ -26,6 +26,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
 
 	"github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
 	"github.com/crossplane-contrib/provider-aws/pkg/utils/pointer"
@@ -643,6 +644,77 @@ func TestGenerateUpdateNodeGroupInput(t *testing.T) {
 				},
 			},
 		},
+		"DiffTaints": {
+			args: args{
+				name: ngName,
+				p: &manualv1alpha1.NodeGroupParameters{
+					ClusterName: clusterName,
+					Taints: []manualv1alpha1.Taint{
+						{
+							Effect: "effect",
+							Key:    ptr.To("toAdd"),
+							Value:  ptr.To("value"),
+						},
+						{
+							Effect: "effect",
+							Key:    ptr.To("toChange"),
+							Value:  ptr.To("newValue"),
+						},
+						{
+							Effect: "effect",
+							Key:    ptr.To("toKeep"),
+							Value:  ptr.To("value"),
+						},
+					},
+				},
+				n: &ekstypes.Nodegroup{
+					NodegroupName: &ngName,
+					ClusterName:   &clusterName,
+					Taints: []ekstypes.Taint{
+						{
+							Effect: "effect",
+							Key:    ptr.To("toKeep"),
+							Value:  ptr.To("value"),
+						},
+						{
+							Effect: "effect",
+							Key:    ptr.To("toRemove"),
+							Value:  ptr.To("value"),
+						},
+						{
+							Effect: "effect",
+							Key:    ptr.To("toChange"),
+							Value:  ptr.To("oldValue"),
+						},
+					},
+				},
+			},
+			want: &eks.UpdateNodegroupConfigInput{
+				NodegroupName: &ngName,
+				ClusterName:   &clusterName,
+				Taints: &ekstypes.UpdateTaintsPayload{
+					AddOrUpdateTaints: []ekstypes.Taint{
+						{
+							Effect: "effect",
+							Key:    ptr.To("toAdd"),
+							Value:  ptr.To("value"),
+						},
+						{
+							Effect: "effect",
+							Key:    ptr.To("toChange"),
+							Value:  ptr.To("newValue"),
+						},
+					},
+					RemoveTaints: []ekstypes.Taint{
+						{
+							Effect: "effect",
+							Key:    ptr.To("toRemove"),
+							Value:  ptr.To("value"),
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for name, tc := range cases {