From f018a49457c46823c2607938f4f45e8c35240250 Mon Sep 17 00:00:00 2001
From: Lucas Pinheiro <lpinheiro@renderedtext.com>
Date: Tue, 23 Jan 2024 08:50:54 -0300
Subject: [PATCH] fix: agent type name update not evicting old agent type from
 registry (#9)

---
 pkg/agenttypes/registry.go      | 22 +++++++++++++++++++---
 pkg/agenttypes/registry_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/pkg/agenttypes/registry.go b/pkg/agenttypes/registry.go
index 71bf0fd..13e1f10 100644
--- a/pkg/agenttypes/registry.go
+++ b/pkg/agenttypes/registry.go
@@ -52,14 +52,30 @@ func (r *Registry) OnUpdate(oldObj, newObj interface{}) {
 		return
 	}
 
-	agentType, err := parseAgentType(newSecret)
+	newAgentType, err := parseAgentType(newSecret)
 	if err != nil {
 		klog.Errorf("Error when parsing agent type: %v", err)
 		return
 	}
 
-	klog.Infof("Agent type updated: %s", agentType.AgentTypeName)
-	r.agentTypes[agentType.AgentTypeName] = agentType
+	oldAgentType, err := parseAgentType(oldSecret)
+	if err != nil {
+		klog.Errorf("Error when parsing agent type: %v", err)
+		return
+	}
+
+	// If the agent type name remained the same, we only need to update it.
+	if oldAgentType.AgentTypeName == newAgentType.AgentTypeName {
+		klog.Infof("Agent type updated: %s", newAgentType.AgentTypeName)
+		r.agentTypes[newAgentType.AgentTypeName] = newAgentType
+		return
+	}
+
+	// Otherwise, we need to delete the old one before adding the new one
+	klog.Infof("Agent type deleted: %s", oldAgentType.AgentTypeName)
+	delete(r.agentTypes, oldAgentType.AgentTypeName)
+	klog.Infof("Agent type added: %s", newAgentType.AgentTypeName)
+	r.agentTypes[newAgentType.AgentTypeName] = newAgentType
 }
 
 func (r *Registry) OnDelete(obj interface{}) {
diff --git a/pkg/agenttypes/registry_test.go b/pkg/agenttypes/registry_test.go
index bd9888b..dfb2cff 100644
--- a/pkg/agenttypes/registry_test.go
+++ b/pkg/agenttypes/registry_test.go
@@ -119,6 +119,36 @@ func Test__Registry(t *testing.T) {
 		require.Equal(t, newToken, agentType.RegistrationToken)
 	})
 
+	t.Run("on update -> old is deleted, new is added", func(t *testing.T) {
+		oldAgentTypeName := randAgentTypeName()
+		secretName := randSecretName()
+		oldToken := randomString(t)
+		require.Nil(t, r.Get(oldAgentTypeName))
+
+		old := &v1.Secret{
+			ObjectMeta: metav1.ObjectMeta{Name: secretName, ResourceVersion: "1"},
+			Data: map[string][]byte{
+				"agentTypeName":     []byte(oldAgentTypeName),
+				"registrationToken": []byte(oldToken),
+			},
+		}
+
+		// agent type is added
+		r.OnAdd(old, false)
+		require.NotNil(t, r.Get(oldAgentTypeName))
+
+		// agent type is updated with different name
+		new := old.DeepCopy()
+		newAgentTypeName := randAgentTypeName()
+		new.Data["agentTypeName"] = []byte(newAgentTypeName)
+		new.ObjectMeta.ResourceVersion = "2"
+		r.OnUpdate(old, new)
+
+		// old is deleted, new is added
+		require.Nil(t, r.Get(oldAgentTypeName))
+		require.NotNil(t, r.Get(newAgentTypeName))
+	})
+
 	t.Run("on delete -> agent type is deleted", func(t *testing.T) {
 		agentTypeName := randAgentTypeName()
 		secretName := randSecretName()