update DockerMachineTemplate instead of replacing it

templates/docker/Dockerfile

@@ -36,11 +36,7 @@ FROM $BUILD_BASE AS builder
 ARG REPO=https://github.com/canonical/k8s-snap
 ARG BRANCH=main
 
-## Override the Kubernetes version from the branch.
-## Note(ben): We only have 1.30-release branches for k8s-snap right now.
-## For the rollout upgrades, we need to have a different minor version.
-## This is a temporary solution until we have a 1.31 release branch.
-ARG KUBERNETES_VERSION_OVERRIDE=""
+ARG KUBERNETES_VERSION=""
 
 ## NOTE(neoaggelos): install dependencies needed to build the tools
 ## !!!IMPORTANT!!! Keep up to date with "snapcraft.yaml:parts.build-deps.build-packages"
@@ -92,12 +88,8 @@ RUN /src/k8s-snap/build-scripts/build-component.sh helm
 
 ## kubernetes build
 FROM builder AS build-kubernetes
-ENV KUBERNETES_VERSION_OVERRIDE=${KUBERNETES_VERSION_OVERRIDE}
-RUN if [ -n "$KUBERNETES_VERSION_OVERRIDE" ]; then \
-      echo "Overwriting Kubernetes version with $KUBERNETES_VERSION_OVERRIDE"; \
-      echo "$KUBERNETES_VERSION_OVERRIDE" > /src/k8s-snap/build-scripts/components/kubernetes/version; \
-      cat /src/k8s-snap/build-scripts/components/kubernetes/version; \
-    fi
+ENV KUBERNETES_VERSION=${KUBERNETES_VERSION}
+RUN [ -n "$KUBERNETES_VERSION" ] && echo "$KUBERNETES_VERSION" > /src/k8s-snap/build-scripts/components/kubernetes/version
 RUN /src/k8s-snap/build-scripts/build-component.sh kubernetes
 
 ## runc build

test/e2e/cluster_upgrade.go

@@ -152,12 +152,18 @@ func ClusterUpgradeSpec(ctx context.Context, inputGetter func() ClusterUpgradeSp
 			WaitForMachineDeployments:    input.E2EConfig.GetIntervals(specName, "wait-worker-nodes"),
 		}, result)
 
+		By("Upgrading the DockerMachineTemplate")
+		UpgradeDockerMachineTemplateAndWaitForUpgrade(ctx, UpgradeDockerMachineTemplateAndWaitForUpgradeInput{
+			ClusterProxy: input.BootstrapClusterProxy,
+			ControlPlane: result.ControlPlane,
+			CustomImage:  "k8s-snap:dev-1.30",
+		})
+
 		By("Upgrading the Kubernetes control-plane")
 		UpgradeControlPlaneAndWaitForUpgrade(ctx, UpgradeControlPlaneAndWaitForUpgradeInput{
 			ClusterProxy:                input.BootstrapClusterProxy,
 			Cluster:                     result.Cluster,
 			ControlPlane:                result.ControlPlane,
-			UpgradeMachineTemplate:      ptr.To(fmt.Sprintf("%s-control-plane-1.30", clusterName)),
 			KubernetesUpgradeVersion:    input.E2EConfig.GetVariable(KubernetesVersionUpgradeTo),
 			WaitForMachinesToBeUpgraded: input.E2EConfig.GetIntervals(specName, "wait-machine-upgrade"),
 		})

test/e2e/data/infrastructure-docker/cluster-template-upgrades.yaml

@@ -60,17 +60,6 @@ spec:
     spec:
       customImage: k8s-snap:dev-1.29
 
-# After upgrade template for the machine deployment
----
-apiVersion: infrastructure.cluster.x-k8s.io/v1beta1
-kind: DockerMachineTemplate
-metadata:
-  name: ${CLUSTER_NAME}-control-plane-1.30
-  namespace: ${NAMESPACE}
-spec:
-  template:
-    spec:
-      customImage: k8s-snap:dev-1.30
 ---
 apiVersion: cluster.x-k8s.io/v1beta1
 kind: MachineDeployment

test/e2e/helpers.go

@@ -37,6 +37,7 @@ import (
 	expv1 "sigs.k8s.io/cluster-api/exp/api/v1beta1"
 	"sigs.k8s.io/cluster-api/test/framework"
 	"sigs.k8s.io/cluster-api/test/framework/clusterctl"
+	dockerv1beta1 "sigs.k8s.io/cluster-api/test/infrastructure/docker/api/v1beta1"
 	"sigs.k8s.io/cluster-api/util/patch"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
@@ -552,13 +553,40 @@ func WaitForControlPlaneAndMachinesReady(ctx context.Context, input WaitForContr
 	})
 }
 
+// UpgradeDockerMachineAndWaitForUpgradeInput is the input type for UpgradeDockerMachineAndWaitForUpgrade.
+type UpgradeDockerMachineTemplateAndWaitForUpgradeInput struct {
+	ClusterProxy framework.ClusterProxy
+	ControlPlane *controlplanev1.CK8sControlPlane
+	CustomImage  string
+}
+
+// UpgradeDockerMachineTemplateAndWaitForUpgrade upgrades a DockerMachineTemplate custom image and waits for it to be upgraded.
+func UpgradeDockerMachineTemplateAndWaitForUpgrade(ctx context.Context, input UpgradeDockerMachineTemplateAndWaitForUpgradeInput) {
+	Byf("Patching the DockerMachineTemplate image to use the updated custom image")
+	mgmtClient := input.ClusterProxy.GetClient()
+
+	dockerMachineTemplate := &dockerv1beta1.DockerMachineTemplate{}
+	err := mgmtClient.Get(ctx, client.ObjectKey{Name: input.ControlPlane.Spec.MachineTemplate.InfrastructureRef.Name, Namespace: input.ControlPlane.Namespace}, dockerMachineTemplate)
+	Expect(err).ToNot(HaveOccurred())
+
+	patchHelperDocker, err := patch.NewHelper(dockerMachineTemplate, mgmtClient)
+	Expect(err).ToNot(HaveOccurred())
+	dockerMachineTemplate.Spec.Template.Spec.CustomImage = input.CustomImage
+	Eventually(func() error {
+		err := patchHelperDocker.Patch(ctx, dockerMachineTemplate)
+		if err != nil {
+			Byf("Failed to patch the DockerMachineTemplate: %v", err)
+		}
+		return err
+	}, retryableOperationTimeout, retryableOperationInterval).Should(Succeed(), "Failed to patch the DockerMachineTemplate")
+}
+
 // UpgradeControlPlaneAndWaitForUpgradeInput is the input type for UpgradeControlPlaneAndWaitForUpgrade.
 type UpgradeControlPlaneAndWaitForUpgradeInput struct {
 	ClusterProxy                framework.ClusterProxy
 	Cluster                     *clusterv1.Cluster
 	ControlPlane                *controlplanev1.CK8sControlPlane
 	KubernetesUpgradeVersion    string
-	UpgradeMachineTemplate      *string
 	WaitForMachinesToBeUpgraded []interface{}
 }
 
@@ -578,17 +606,6 @@ func UpgradeControlPlaneAndWaitForUpgrade(ctx context.Context, input UpgradeCont
 
 	input.ControlPlane.Spec.Version = input.KubernetesUpgradeVersion
 
-	// Create a new ObjectReference for the infrastructure provider
-	newInfrastructureRef := corev1.ObjectReference{
-		APIVersion: "infrastructure.cluster.x-k8s.io/v1beta1", // Adjust based on your infrastructure API version
-		Kind:       "DockerMachineTemplate",
-		Name:       fmt.Sprintf("%s-control-plane-1.30", input.Cluster.Name),
-		Namespace:  input.ControlPlane.Spec.MachineTemplate.InfrastructureRef.Namespace,
-	}
-
-	// Update the infrastructureRef
-	input.ControlPlane.Spec.MachineTemplate.InfrastructureRef = newInfrastructureRef
-
 	Eventually(func() error {
 		return patchHelper.Patch(ctx, input.ControlPlane)
 	}, retryableOperationTimeout, retryableOperationInterval).Should(Succeed(), "Failed to patch the new kubernetes version to KCP %s", klog.KObj(input.ControlPlane))
@@ -618,16 +635,6 @@ func UpgradeMachineDeploymentsAndWait(ctx context.Context, input framework.Upgra
 
 		oldVersion := deployment.Spec.Template.Spec.Version
 		deployment.Spec.Template.Spec.Version = &input.UpgradeVersion
-		// Create a new ObjectReference for the infrastructure provider
-		newInfrastructureRef := corev1.ObjectReference{
-			APIVersion: "infrastructure.cluster.x-k8s.io/v1beta1",
-			Kind:       "DockerMachineTemplate",
-			Name:       fmt.Sprintf("%s-md-1.30-0", input.Cluster.Name),
-			Namespace:  deployment.Spec.Template.Spec.InfrastructureRef.Namespace,
-		}
-
-		// Update the infrastructureRef
-		deployment.Spec.Template.Spec.InfrastructureRef = newInfrastructureRef
 		Eventually(func() error {
 			return patchHelper.Patch(ctx, deployment)
 		}, retryableOperationTimeout, retryableOperationInterval).Should(Succeed(), "Failed to patch Kubernetes version on MachineDeployment %s", klog.KObj(deployment))