From 210090d34c5aa5c852ea3d0ea6bbc3376a458e66 Mon Sep 17 00:00:00 2001
From: Markus Rudy <webmaster@burgerdev.de>
Date: Thu, 14 Dec 2023 17:29:04 +0100
Subject: [PATCH] fixup! upgrade-agent: allow more than one KubernetesComponent

---
 upgrade-agent/internal/server/server_test.go | 55 ++++++++++++++++++--
 1 file changed, 51 insertions(+), 4 deletions(-)

diff --git a/upgrade-agent/internal/server/server_test.go b/upgrade-agent/internal/server/server_test.go
index f441900cb..a418575a2 100644
--- a/upgrade-agent/internal/server/server_test.go
+++ b/upgrade-agent/internal/server/server_test.go
@@ -56,9 +56,38 @@ func TestVersionVerifier(t *testing.T) {
 }
 
 func TestPrepareUpdate(t *testing.T) {
-	validUpdateRequest := &upgradeproto.ExecuteUpdateRequest{
+	invalidUpgradeRequest := &upgradeproto.ExecuteUpdateRequest{
+		WantedKubernetesVersion: "1337",
+	}
+	slimUpdateRequest := &upgradeproto.ExecuteUpdateRequest{
 		WantedKubernetesVersion: "v1.1.1",
 	}
+	oldStyleUpdateRequest := &upgradeproto.ExecuteUpdateRequest{
+		WantedKubernetesVersion: "v1.1.1",
+		KubeadmUrl:              "http://example.com/kubeadm",
+		KubeadmHash:             "sha256:foo",
+	}
+	newStyleUpdateRequest := &upgradeproto.ExecuteUpdateRequest{
+		WantedKubernetesVersion: "v1.1.1",
+		KubernetesComponents: []*components.Component{
+			{
+				Url:         "http://example.com/kubeadm",
+				Hash:        "sha256:foo",
+				InstallPath: "/tmp/kubeadm",
+			},
+		},
+	}
+	combinedStyleUpdateRequest := &upgradeproto.ExecuteUpdateRequest{
+		WantedKubernetesVersion: "v1.1.1",
+		KubeadmUrl:              "http://example.com/kubeadm",
+		KubeadmHash:             "sha256:foo",
+		KubernetesComponents: []*components.Component{
+			{
+				Url:         "data:application/octet-stream,foo",
+				InstallPath: "/tmp/foo",
+			},
+		},
+	}
 	testCases := map[string]struct {
 		installer     osInstaller
 		updateRequest *upgradeproto.ExecuteUpdateRequest
@@ -66,16 +95,34 @@ func TestPrepareUpdate(t *testing.T) {
 	}{
 		"works": {
 			installer:     stubOsInstaller{},
-			updateRequest: validUpdateRequest,
+			updateRequest: slimUpdateRequest,
 		},
 		"invalid version string": {
 			installer:     stubOsInstaller{},
-			updateRequest: &upgradeproto.ExecuteUpdateRequest{WantedKubernetesVersion: "1337"},
+			updateRequest: invalidUpgradeRequest,
 			wantErr:       true,
 		},
 		"install error": {
 			installer:     stubOsInstaller{InstallErr: fmt.Errorf("install error")},
-			updateRequest: validUpdateRequest,
+			updateRequest: oldStyleUpdateRequest,
+			wantErr:       true,
+		},
+		"new style works": {
+			installer:     stubOsInstaller{},
+			updateRequest: newStyleUpdateRequest,
+		},
+		"new style install error": {
+			installer:     stubOsInstaller{InstallErr: fmt.Errorf("install error")},
+			updateRequest: newStyleUpdateRequest,
+			wantErr:       true,
+		},
+		"combined style works": {
+			installer:     stubOsInstaller{},
+			updateRequest: combinedStyleUpdateRequest,
+		},
+		"combined style install error": {
+			installer:     stubOsInstaller{InstallErr: fmt.Errorf("install error")},
+			updateRequest: combinedStyleUpdateRequest,
 			wantErr:       true,
 		},
 	}