Add some tests around channel adding needs-update annotation

channels/pkg/channels/BUILD.bazel

@@ -46,6 +46,7 @@ go_test(
         "//vendor/github.com/stretchr/testify/require:go_default_library",
         "//vendor/k8s.io/api/core/v1:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
+        "//vendor/k8s.io/apimachinery/pkg/runtime:go_default_library",
         "//vendor/k8s.io/client-go/kubernetes/fake:go_default_library",
     ],
 )

channels/pkg/channels/addon.go

@@ -173,14 +173,9 @@ func (a *Addon) EnsureUpdated(ctx context.Context, k8sClient kubernetes.Interfac
 			return nil, fmt.Errorf("error applying update from %q: %v", manifestURL, err)
 		}
 
-		if required.ExistingVersion != nil {
+		if err := a.AddNeedsUpdateLabel(ctx, k8sClient, required); err != nil {
-			if a.Spec.NeedsRollingUpdate != "" {
-				err = a.AddNeedsUpdateLabel(ctx, k8sClient)
-				if err != nil {
 			return nil, fmt.Errorf("error adding needs-update label: %v", err)
 		}
-			}
-		}
 
 		channel := a.buildChannel()
 		err = channel.SetInstalledVersion(ctx, k8sClient, a.ChannelVersion())
@@ -197,7 +192,19 @@ func (a *Addon) EnsureUpdated(ctx context.Context, k8sClient kubernetes.Interfac
 	return required, nil
 }
 
-func (a *Addon) AddNeedsUpdateLabel(ctx context.Context, k8sClient kubernetes.Interface) error {
+func (a *Addon) AddNeedsUpdateLabel(ctx context.Context, k8sClient kubernetes.Interface, required *AddonUpdate) error {
+	if required.ExistingVersion != nil {
+		if a.Spec.NeedsRollingUpdate != "" {
+			err := a.patchNeedsUpdateLabel(ctx, k8sClient)
+			if err != nil {
+				return fmt.Errorf("error patching needs-update label: %v", err)
+			}
+		}
+	}
+	return nil
+}
+
+func (a *Addon) patchNeedsUpdateLabel(ctx context.Context, k8sClient kubernetes.Interface) error {
 	klog.Infof("addon %v wants to update %v nodes", a.Name, a.Spec.NeedsRollingUpdate)
 	selector := ""
 	switch a.Spec.NeedsRollingUpdate {

channels/pkg/channels/addons_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
 	fakekubernetes "k8s.io/client-go/kubernetes/fake"
 	"k8s.io/kops/channels/pkg/api"
 
@@ -259,6 +260,185 @@ func Test_GetRequiredUpdates(t *testing.T) {
 	}
 }
 
+func Test_NeedsRollingUpdate(t *testing.T) {
+	grid := []struct {
+		newAddon            *Addon
+		originalAnnotations map[string]string
+		updateRequired      bool
+		installRequired     bool
+		expectedNodeUpdates int
+	}{
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					ManifestHash:       "originalHash",
+					Version:            fi.String("1"),
+					NeedsRollingUpdate: "all",
+				},
+			},
+		},
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					ManifestHash:       "originalHash",
+					Version:            fi.String("2"),
+					NeedsRollingUpdate: "all",
+				},
+			},
+			updateRequired:      true,
+			expectedNodeUpdates: 2,
+		},
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					Version:            fi.String("1"),
+					ManifestHash:       "newHash",
+					NeedsRollingUpdate: "all",
+				},
+			},
+			updateRequired:      true,
+			expectedNodeUpdates: 2,
+		},
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					Version:            fi.String("1"),
+					ManifestHash:       "newHash",
+					NeedsRollingUpdate: "worker",
+				},
+			},
+			updateRequired:      true,
+			expectedNodeUpdates: 1,
+		},
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					Version:            fi.String("1"),
+					ManifestHash:       "newHash",
+					NeedsRollingUpdate: "control-plane",
+				},
+			},
+			updateRequired:      true,
+			expectedNodeUpdates: 1,
+		},
+		{
+			newAddon: &Addon{
+				Name: "test",
+				Spec: &api.AddonSpec{
+					Name:               fi.String("test"),
+					Version:            fi.String("1"),
+					ManifestHash:       "newHash",
+					NeedsRollingUpdate: "all",
+				},
+			},
+			originalAnnotations: map[string]string{
+				"addons.k8s.io/placeholder": "{\"version\":\"1\",\"manifestHash\":\"originalHash\"}",
+			},
+			installRequired:     true,
+			expectedNodeUpdates: 0,
+		},
+	}
+
+	for _, g := range grid {
+		ctx := context.Background()
+
+		annotations := map[string]string{
+			"addons.k8s.io/test": "{\"version\":\"1\",\"manifestHash\":\"originalHash\"}",
+		}
+		if len(g.originalAnnotations) > 0 {
+			annotations = g.originalAnnotations
+		}
+
+		objects := []runtime.Object{
+			&corev1.Namespace{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:        "kube-system",
+					Annotations: annotations,
+				},
+			},
+			&corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "cp",
+					Labels: map[string]string{
+						"node-role.kubernetes.io/master": "",
+					},
+				},
+			},
+			&corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "node",
+					Labels: map[string]string{
+						"node-role.kubernetes.io/node": "",
+					},
+				},
+			},
+		}
+		fakek8s := fakekubernetes.NewSimpleClientset(objects...)
+		fakecm := fakecertmanager.NewSimpleClientset()
+
+		addon := g.newAddon
+		required, err := addon.GetRequiredUpdates(ctx, fakek8s, fakecm)
+		if err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+
+		if !g.updateRequired && !g.installRequired {
+			if required == nil {
+				continue
+			} else {
+				t.Fatalf("did not expect update, but required was not nil")
+			}
+		}
+
+		if required == nil {
+			t.Fatalf("expected required update, got nil")
+		}
+
+		if required.NewVersion == nil {
+			t.Errorf("updating or installing addon, but NewVersion was nil")
+		}
+
+		if required.ExistingVersion != nil {
+			if g.installRequired {
+				t.Errorf("new install of addon, but ExistingVersion was not nil")
+			}
+		} else {
+			if g.updateRequired {
+				t.Errorf("update of addon, but ExistingVersion was nil")
+			}
+		}
+
+		if err := addon.AddNeedsUpdateLabel(ctx, fakek8s, required); err != nil {
+			t.Errorf("unexpected error: %v", err)
+		}
+
+		nodes, _ := fakek8s.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
+		nodeUpdates := 0
+
+		for _, node := range nodes.Items {
+			if _, exists := node.Annotations["kops.k8s.io/needs-update"]; exists {
+				nodeUpdates++
+			}
+		}
+
+		if nodeUpdates != g.expectedNodeUpdates {
+			t.Errorf("expected %d node updates, but got %d", g.expectedNodeUpdates, nodeUpdates)
+		}
+
+	}
+
+}
+
 func Test_InstallPKI(t *testing.T) {
 	ctx := context.Background()
 	kubeSystem := &corev1.Namespace{