Repository: rancher-sandbox/cluster-api-provider-rke2
Pull request #306 (2024-04-22 13:13:32)

Use spec.Version field as primarily source
<!--
Label the PR with the kind of change this for:

kind/feature
kind/bug
kind/documentation
kind/regression
kind/*
-->

**What this PR does / why we need it**:

Use spec.Version field as primarily source of RKE2 version. Currently, it's just a placeholder for satisfying API contracts. When both spec.AgentConfig.Version and spec.Version are specified, spec.Version take precedence.

<!-- Enter a description of the change and why this change is needed -->

**Which issue(s) this PR fixes** *(optional, in `fixes #<issue number>(, fixes #<issue_number>, ...)` format, will close the issue(s) when PR gets merged)*:
Fixes #

**Special notes for your reviewer**:

**Checklist**:
<!-- Put an "X" character inside the brackets of each completed task. Some may be optional depending on the PR in which case these can be deleted -->

- [ ] squashed commits into logical changes
- [ ] includes documentation
- [ ] adds unit tests
- [ ] adds or updates e2e tests


alexander-demicev: @Danil-Grigorev this should stay a valid option, I don't want to deprecate a feature for users without notice. We can make it only possible place to specify version in future releases. 

Patch:
diff --git a/bootstrap/api/v1beta1/rke2config_types.go b/bootstrap/api/v1beta1/rke2config_types.go
index 223b978..197e5b7 100644
--- a/bootstrap/api/v1beta1/rke2config_types.go
+++ b/bootstrap/api/v1beta1/rke2config_types.go
@@ -148,6 +148,7 @@ type RKE2AgentConfig struct {
 
 	// Version specifies the rke2 version.
 	// This field will be deprecated in newer versions of the API and RKE2ControlPlaneSpec.Version will be used instead.
+	// +kubebuilder:validation:Pattern="v(\\d\\.\\d{2}\\.\\d)\\+rke2r\\d"
 	//+optional
 	Version string `json:"version,omitempty"`
 
diff --git a/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configs.yaml b/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configs.yaml
index f109319..a1034f3 100644
--- a/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configs.yaml
+++ b/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configs.yaml
@@ -892,6 +892,7 @@ spec:
                     description: |-
                       Version specifies the rke2 version.
                       This field will be deprecated in newer versions of the API and RKE2ControlPlaneSpec.Version will be used instead.
+                    pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                     type: string
                 type: object
               files:
diff --git a/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configtemplates.yaml b/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configtemplates.yaml
index bd07252..28e96cc 100644
--- a/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configtemplates.yaml
+++ b/bootstrap/config/crd/bases/bootstrap.cluster.x-k8s.io_rke2configtemplates.yaml
@@ -875,6 +875,7 @@ spec:
                             description: |-
                               Version specifies the rke2 version.
                               This field will be deprecated in newer versions of the API and RKE2ControlPlaneSpec.Version will be used instead.
+                            pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                             type: string
                         type: object
                       files:
diff --git a/bootstrap/internal/controllers/rke2config_controller.go b/bootstrap/internal/controllers/rke2config_controller.go
index d65ba5b..e4e84e7 100644
--- a/bootstrap/internal/controllers/rke2config_controller.go
+++ b/bootstrap/internal/controllers/rke2config_controller.go
@@ -278,6 +278,14 @@ type Scope struct {
 	ControlPlane         *controlplanev1.RKE2ControlPlane
 }
 
+func (s *Scope) getDesiredVersion() string {
+	if s.Machine.Spec.Version != nil && bsutil.IsRKE2Version(*s.Machine.Spec.Version) {
+		return *s.Machine.Spec.Version
+	}
+
+	return s.Config.Spec.AgentConfig.Version
+}
+
 // SetupWithManager sets up the controller with the Manager.
 func (r *RKE2ConfigReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	if r.RKE2InitLock == nil {
@@ -363,6 +371,7 @@ func (r *RKE2ConfigReconciler) handleClusterNotInitialized(ctx context.Context,
 			AgentConfig:          scope.Config.Spec.AgentConfig,
 			Ctx:                  ctx,
 			Client:               r.Client,
+			Version:              scope.getDesiredVersion(),
 		})
 	if err != nil {
 		return ctrl.Result{}, err
@@ -417,7 +426,7 @@ func (r *RKE2ConfigReconciler) handleClusterNotInitialized(ctx context.Context,
 			PreRKE2Commands:         scope.Config.Spec.PreRKE2Commands,
 			PostRKE2Commands:        scope.Config.Spec.PostRKE2Commands,
 			ConfigFile:              initConfigFile,
-			RKE2Version:             scope.Config.Spec.AgentConfig.Version,
+			RKE2Version:             scope.getDesiredVersion(),
 			WriteFiles:              files,
 			NTPServers:              ntpServers,
 			AdditionalCloudInit:     scope.Config.Spec.AgentConfig.AdditionalUserData.Config,
@@ -567,6 +576,7 @@ func (r *RKE2ConfigReconciler) joinControlplane(ctx context.Context, scope *Scop
 			AgentConfig:          scope.Config.Spec.AgentConfig,
 			Ctx:                  ctx,
 			Client:               r.Client,
+			Version:              scope.getDesiredVersion(),
 		},
 	)
 	if err != nil {
@@ -627,7 +637,7 @@ func (r *RKE2ConfigReconciler) joinControlplane(ctx context.Context, scope *Scop
 			PreRKE2Commands:     scope.Config.Spec.PreRKE2Commands,
 			PostRKE2Commands:    scope.Config.Spec.PostRKE2Commands,
 			ConfigFile:          initConfigFile,
-			RKE2Version:         scope.Config.Spec.AgentConfig.Version,
+			RKE2Version:         scope.getDesiredVersion(),
 			WriteFiles:          files,
 			NTPServers:          ntpServers,
 			AdditionalCloudInit: scope.Config.Spec.AgentConfig.AdditionalUserData.Config,
@@ -700,6 +710,7 @@ func (r *RKE2ConfigReconciler) joinWorker(ctx context.Context, scope *Scope) (re
 			Client:                 r.Client,
 			CloudProviderName:      scope.ControlPlane.Spec.ServerConfig.CloudProviderName,
 			CloudProviderConfigMap: scope.ControlPlane.Spec.ServerConfig.CloudProviderConfigMap,
+			Version:                scope.getDesiredVersion(),
 		})
 	if err != nil {
 		return ctrl.Result{}, err
@@ -738,7 +749,7 @@ func (r *RKE2ConfigReconciler) joinWorker(ctx context.Context, scope *Scope) (re
 		CISEnabled:              scope.Config.Spec.AgentConfig.CISProfile != "",
 		PostRKE2Commands:        scope.Config.Spec.PostRKE2Commands,
 		ConfigFile:              wkJoinConfigFile,
-		RKE2Version:             scope.Config.Spec.AgentConfig.Version,
+		RKE2Version:             scope.getDesiredVersion(),
 		WriteFiles:              files,
 		NTPServers:              ntpServers,
 		AdditionalCloudInit:     scope.Config.Spec.AgentConfig.AdditionalUserData.Config,
diff --git a/controlplane/api/v1beta1/rke2controlplane_types.go b/controlplane/api/v1beta1/rke2controlplane_types.go
index 4635d68..d9e12f5 100644
--- a/controlplane/api/v1beta1/rke2controlplane_types.go
+++ b/controlplane/api/v1beta1/rke2controlplane_types.go
@@ -17,6 +17,8 @@ limitations under the License.
 package v1beta1
 
 import (
+	"regexp"
+
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -47,9 +49,8 @@ type RKE2ControlPlaneSpec struct {
 	Replicas *int32 `json:"replicas,omitempty"`
 
 	// Version defines the desired Kubernetes version.
-	// This is only a placeholder for now, and the RKE2ConfigSpec.AgentConfig.Version field should be used instead.
-	// In future iterations, this field overrides the RKE2 Version specificied in RKE2ConfigSpec.AgentConfig.Version
-	// which will be deprecated in newer versions of the API.
+	// This field takes precedence over RKE2ConfigSpec.AgentConfig.Version (which is deprecated).
+	// +kubebuilder:validation:Pattern="v(\\d\\.\\d{2}\\.\\d)\\+rke2r\\d"
 	// +optional
 	Version string `json:"version"`
 
@@ -436,3 +437,21 @@ func (r *RKE2ControlPlane) GetConditions() clusterv1.Conditions {
 func (r *RKE2ControlPlane) SetConditions(conditions clusterv1.Conditions) {
 	r.Status.Conditions = conditions
 }
+
+// GetDesiredVersion returns the desired version of the RKE2ControlPlane using Spec.Version field as a default field.
+func (r *RKE2ControlPlane) GetDesiredVersion() string {
+	if r.Spec.Version != "" && isRKE2Version(r.Spec.Version) {
+		return r.Spec.Version
+	}
+
+	return r.Spec.AgentConfig.Version
+}
+
+// isRKE2Version checks if a string is an RKE2 version.
+func isRKE2Version(rke2Version string) bool {
+	regexStr := "v(\\d\\.\\d{2}\\.\\d)\\+rke2r\\d"
+
+	regex, _ := regexp.Compile(regexStr)
+
+	return regex.MatchString(rke2Version)
+}
diff --git a/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanes.yaml b/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanes.yaml
index cc7e943..943f3c5 100644
--- a/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanes.yaml
+++ b/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanes.yaml
@@ -1537,6 +1537,7 @@ spec:
                     description: |-
                       Version specifies the rke2 version.
                       This field will be deprecated in newer versions of the API and RKE2ControlPlaneSpec.Version will be used instead.
+                    pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                     type: string
                 type: object
               files:
@@ -2449,9 +2450,8 @@ spec:
               version:
                 description: |-
                   Version defines the desired Kubernetes version.
-                  This is only a placeholder for now, and the RKE2ConfigSpec.AgentConfig.Version field should be used instead.
-                  In future iterations, this field overrides the RKE2 Version specificied in RKE2ConfigSpec.AgentConfig.Version
-                  which will be deprecated in newer versions of the API.
+                  This field takes precedence over RKE2ConfigSpec.AgentConfig.Version (which is deprecated).
+                pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                 type: string
             required:
             - infrastructureRef
diff --git a/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanetemplates.yaml b/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanetemplates.yaml
index 8820c9a..1e283df 100644
--- a/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanetemplates.yaml
+++ b/controlplane/config/crd/bases/controlplane.cluster.x-k8s.io_rke2controlplanetemplates.yaml
@@ -368,6 +368,7 @@ spec:
                             description: |-
                               Version specifies the rke2 version.
                               This field will be deprecated in newer versions of the API and RKE2ControlPlaneSpec.Version will be used instead.
+                            pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                             type: string
                         type: object
                       files:
@@ -1300,9 +1301,8 @@ spec:
                       version:
                         description: |-
                           Version defines the desired Kubernetes version.
-                          This is only a placeholder for now, and the RKE2ConfigSpec.AgentConfig.Version field should be used instead.
-                          In future iterations, this field overrides the RKE2 Version specificied in RKE2ConfigSpec.AgentConfig.Version
-                          which will be deprecated in newer versions of the API.
+                          This field takes precedence over RKE2ConfigSpec.AgentConfig.Version (which is deprecated).
+                        pattern: v(\d\.\d{2}\.\d)\+rke2r\d
                         type: string
                     required:
                     - infrastructureRef
diff --git a/controlplane/internal/controllers/rke2controlplane_controller.go b/controlplane/internal/controllers/rke2controlplane_controller.go
index f209c41..ffd5bb3 100644
--- a/controlplane/internal/controllers/rke2controlplane_controller.go
+++ b/controlplane/internal/controllers/rke2controlplane_controller.go
@@ -609,9 +609,9 @@ func (r *RKE2ControlPlaneReconciler) reconcileEtcdMembers(ctx context.Context, c
 		return errors.Wrap(err, "cannot get remote client to workload cluster")
 	}
 
-	parsedVersion, err := semver.ParseTolerant(controlPlane.RCP.Spec.AgentConfig.Version)
+	parsedVersion, err := semver.ParseTolerant(controlPlane.RCP.GetDesiredVersion())
 	if err != nil {
-		return errors.Wrapf(err, "failed to parse kubernetes version %q", controlPlane.RCP.Spec.AgentConfig.Version)
+		return errors.Wrapf(err, "failed to parse kubernetes version %q", controlPlane.RCP.GetDesiredVersion())
 	}
 
 	removedMembers, err := workloadCluster.ReconcileEtcdMembers(ctx, nodeNames, parsedVersion)
diff --git a/controlplane/internal/controllers/scale.go b/controlplane/internal/controllers/scale.go
index 1781096..e01ce96 100644
--- a/controlplane/internal/controllers/scale.go
+++ b/controlplane/internal/controllers/scale.go
@@ -19,7 +19,6 @@ package controllers
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"strings"
 	"time"
 
@@ -42,7 +41,6 @@ import (
 	bootstrapv1 "github.com/rancher-sandbox/cluster-api-provider-rke2/bootstrap/api/v1beta1"
 	controlplanev1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1beta1"
 	rke2 "github.com/rancher-sandbox/cluster-api-provider-rke2/pkg/rke2"
-	bsutil "github.com/rancher-sandbox/cluster-api-provider-rke2/pkg/util"
 )
 
 func (r *RKE2ControlPlaneReconciler) initializeControlPlane(
@@ -413,14 +411,11 @@ func (r *RKE2ControlPlaneReconciler) generateMachine(
 	bootstrapRef *corev1.ObjectReference,
 	failureDomain *string,
 ) error {
-	newVersion, err := bsutil.Rke2ToKubeVersion(rcp.Spec.AgentConfig.Version)
-	if err != nil {
-		return fmt.Errorf("failed to convert rke2 version to kubernetes version: %w", err)
-	}
-
 	logger := log.FromContext(ctx)
 
-	logger.Info("Version checking...", "rke2-version", rcp.Spec.AgentConfig.Version, "machine-version: ", newVersion)
+	rke2Version := rcp.GetDesiredVersion()
+
+	logger.Info("Version checking...", "rke2-version", rke2Version)
 
 	machine := &clusterv1.Machine{
 		ObjectMeta: metav1.ObjectMeta{
@@ -433,7 +428,7 @@ func (r *RKE2ControlPlaneReconciler) generateMachine(
 		},
 		Spec: clusterv1.MachineSpec{
 			ClusterName:       cluster.Name,
-			Version:           &newVersion,
+			Version:           &rke2Version,
 			InfrastructureRef: *infraRef,
 			Bootstrap: clusterv1.Bootstrap{
 				ConfigRef: bootstrapRef,
diff --git a/pkg/rke2/config.go b/pkg/rke2/config.go
index 573b624..71af779 100644
--- a/pkg/rke2/config.go
+++ b/pkg/rke2/config.go
@@ -155,6 +155,7 @@ type ServerConfigOpts struct {
 	AgentConfig          bootstrapv1.RKE2AgentConfig
 	Ctx                  context.Context
 	Client               client.Client
+	Version              string
 }
 
 func newRKE2ServerConfig(opts ServerConfigOpts) (*rke2ServerConfig, []bootstrapv1.File, error) { // nolint:gocyclo
@@ -403,6 +404,7 @@ type AgentConfigOpts struct {
 	Client                 client.Client
 	CloudProviderName      string
 	CloudProviderConfigMap *corev1.ObjectReference
+	Version                string
 }
 
 func newRKE2AgentConfig(opts AgentConfigOpts) (*rke2AgentConfig, []bootstrapv1.File, error) {
@@ -411,8 +413,8 @@ func newRKE2AgentConfig(opts AgentConfigOpts) (*rke2AgentConfig, []bootstrapv1.F
 	rke2AgentConfig.ContainerRuntimeEndpoint = opts.AgentConfig.ContainerRuntimeEndpoint
 
 	if opts.AgentConfig.CISProfile != "" {
-		if !bsutil.ProfileCompliant(opts.AgentConfig.CISProfile, opts.AgentConfig.Version) {
-			return nil, nil, fmt.Errorf("profile %q is not supported for version %q", opts.AgentConfig.CISProfile, opts.AgentConfig.Version)
+		if !bsutil.ProfileCompliant(opts.AgentConfig.CISProfile, opts.Version) {
+			return nil, nil, fmt.Errorf("profile %q is not supported for version %q", opts.AgentConfig.CISProfile, opts.Version)
 		}
 
 		files = append(files, bootstrapv1.File{
@@ -550,6 +552,7 @@ func GenerateInitControlPlaneConfig(opts ServerConfigOpts) (*rke2ServerConfig, [
 		Client:      opts.Client,
 		Ctx:         opts.Ctx,
 		Token:       opts.Token,
+		Version:     opts.Version,
 	})
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to generate rke2 agent config: %w", err)
@@ -581,6 +584,7 @@ func GenerateJoinControlPlaneConfig(opts ServerConfigOpts) (*rke2ServerConfig, [
 		Ctx:         opts.Ctx,
 		ServerURL:   opts.ServerURL,
 		Token:       opts.Token,
+		Version:     opts.Version,
 	})
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to generate rke2 agent config: %w", err)
diff --git a/pkg/rke2/control_plane.go b/pkg/rke2/control_plane.go
index 7e90e7e..ea35d95 100644
--- a/pkg/rke2/control_plane.go
+++ b/pkg/rke2/control_plane.go
@@ -116,7 +116,9 @@ func (c *ControlPlane) FailureDomains() clusterv1.FailureDomains {
 
 // Version returns the RKE2ControlPlane's version.
 func (c *ControlPlane) Version() *string {
-	return &c.RCP.Spec.AgentConfig.Version
+	version := c.RCP.GetDesiredVersion()
+
+	return &version
 }
 
 // InfrastructureRef returns the RKE2ControlPlane's infrastructure template.
diff --git a/pkg/rke2/machine_filters.go b/pkg/rke2/machine_filters.go
index 819235c..28c3b8f 100644
--- a/pkg/rke2/machine_filters.go
+++ b/pkg/rke2/machine_filters.go
@@ -22,7 +22,7 @@ func matchesRCPConfiguration(
 	rcp *controlplanev1.RKE2ControlPlane,
 ) func(machine *clusterv1.Machine) bool {
 	return collections.And(
-		matchesKubernetesVersion(rcp.Spec.AgentConfig.Version),
+		matchesKubernetesOrRKE2Version(rcp.GetDesiredVersion()),
 		matchesRKE2BootstrapConfig(machineConfigs, rcp),
 		matchesTemplateClonedFrom(infraConfigs, rcp),
 	)
@@ -119,8 +119,8 @@ func matchesTemplateClonedFrom(infraConfigs map[string]*unstructured.Unstructure
 	}
 }
 
-// matchesKubernetesVersion returns a filter to find all machines that match a given Kubernetes version.
-func matchesKubernetesVersion(kubernetesVersion string) func(*clusterv1.Machine) bool {
+// matchesKubernetesVersion returns a filter to find all machines that match a given Kubernetes or RKE2 version.
+func matchesKubernetesOrRKE2Version(rke2Version string) func(*clusterv1.Machine) bool {
 	return func(machine *clusterv1.Machine) bool {
 		if machine == nil {
 			return false
@@ -130,7 +130,11 @@ func matchesKubernetesVersion(kubernetesVersion string) func(*clusterv1.Machine)
 			return false
 		}
 
-		rcpKubeVersion, err := bsutil.Rke2ToKubeVersion(kubernetesVersion)
+		if bsutil.IsRKE2Version(*machine.Spec.Version) {
+			return bsutil.CompareVersions(*machine.Spec.Version, rke2Version)
+		}
+
+		rcpKubeVersion, err := bsutil.Rke2ToKubeVersion(rke2Version)
 		if err != nil {
 			return true
 		}
diff --git a/pkg/util/util.go b/pkg/util/util.go
index d2b3a88..c49df1d 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -128,6 +128,15 @@ func Rke2ToKubeVersion(rk2Version string) (kubeVersion string, err error) {
 	return kubeVersion, nil
 }
 
+// IsRKE2Version checks if a string is an RKE2 version.
+func IsRKE2Version(rke2Version string) bool {
+	regexStr := "v(\\d\\.\\d{2}\\.\\d)\\+rke2r\\d"
+
+	regex, _ := regexp.Compile(regexStr)
+
+	return regex.MatchString(rke2Version)
+}
+
 // AppendIfNotPresent appends a string to a slice only if the value does not already exist.
 func AppendIfNotPresent(origSlice []string, strItem string) (resultSlice []string) {
 	present := false


Test patch:
diff --git a/pkg/rke2/config_test.go b/pkg/rke2/config_test.go
index f8c6072..614e454 100644
--- a/pkg/rke2/config_test.go
+++ b/pkg/rke2/config_test.go
@@ -295,8 +295,8 @@ var _ = Describe("RKE2 Agent Config", func() {
 					ExtraEnv:      map[string]string{"testenv": "testenv"},
 					ExtraMounts:   map[string]string{"testmount": "testmount"},
 				},
-				Version: "v1.25.2",
 			},
+			Version: "v1.25.2",
 		}
 	})
 
diff --git a/pkg/rke2/machine_filters_test.go b/pkg/rke2/machine_filters_test.go
index 288023b..efb1bfc 100644
--- a/pkg/rke2/machine_filters_test.go
+++ b/pkg/rke2/machine_filters_test.go
@@ -14,6 +14,12 @@ import (
 	controlplanev1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1beta1"
 )
 
+var (
+	k8sMachineVersion  = "v1.24.6"
+	rke2MachineVersion = "v1.24.6+rke2r1"
+	regionEuCentral1   = "eu-central-1"
+)
+
 var rcp = controlplanev1.RKE2ControlPlane{
 	ObjectMeta: v1.ObjectMeta{
 		Name:      "rke2-cluster-control-plane",
@@ -27,18 +33,13 @@ var rcp = controlplanev1.RKE2ControlPlane{
 		},
 		RKE2ConfigSpec: bootstrapv1.RKE2ConfigSpec{
 			AgentConfig: bootstrapv1.RKE2AgentConfig{
-				Version:    "v1.24.6+rke2r1",
+				Version:    rke2MachineVersion,
 				NodeLabels: []string{"hello=world"},
 			},
 		},
 	},
 }
 
-var (
-	machineVersion   = "v1.24.6"
-	regionEuCentral1 = "eu-central-1"
-)
-
 var machine = clusterv1.Machine{
 	ObjectMeta: v1.ObjectMeta{
 		Name:      "machine-test",
@@ -49,7 +50,7 @@ var machine = clusterv1.Machine{
 	},
 	Spec: clusterv1.MachineSpec{
 		ClusterName:   "rke2-cluster",
-		Version:       &machineVersion,
+		Version:       &k8sMachineVersion,
 		FailureDomain: &regionEuCentral1,
 		Bootstrap: clusterv1.Bootstrap{
 			ConfigRef: &corev1.ObjectReference{
@@ -80,7 +81,7 @@ var _ = Describe("matchAgentConfig", func() {
 				},
 				Spec: bootstrapv1.RKE2ConfigSpec{
 					AgentConfig: bootstrapv1.RKE2AgentConfig{
-						Version:    "v1.24.6+rke2r1",
+						Version:    rke2MachineVersion,
 						NodeLabels: []string{"hello=world"},
 					},
 				},
@@ -99,7 +100,15 @@ var _ = Describe("matchAgentConfig", func() {
 var _ = Describe("matching Kubernetes Version", func() {
 	It("should match version", func() {
 		machineCollection := collections.FromMachines(&machine)
-		matches := machineCollection.AnyFilter(matchesKubernetesVersion(rcp.Spec.AgentConfig.Version))
+		matches := machineCollection.AnyFilter(matchesKubernetesOrRKE2Version(rcp.GetDesiredVersion()))
+		Expect(len(matches)).To(Equal(1))
+	})
+
+	It("should match when RKE2 version is set on the machine", func() {
+		machine.Spec.Version = &rke2MachineVersion
+		machineCollection := collections.FromMachines(&machine)
+		matches := machineCollection.AnyFilter(matchesKubernetesOrRKE2Version(rcp.GetDesiredVersion()))
 		Expect(len(matches)).To(Equal(1))
+		machine.Spec.Version = &k8sMachineVersion
 	})
 })
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
index 1c9b544..01a5431 100644
--- a/pkg/util/util_test.go
+++ b/pkg/util/util_test.go
@@ -26,3 +26,16 @@ var _ = Describe(("Testing GetMapKeysAsString"), func() {
 		Expect(keysValid).To(BeTrue())
 	})
 })
+
+var _ = Describe("Testing IsRKE2Version", func() {
+	k8sVersion := "1.24.6"
+	rke2Version := "v1.24.6+rke2r1"
+
+	It("Should return true if string is RKE2 version", func() {
+		Expect(IsRKE2Version(rke2Version)).To(BeTrue())
+	})
+
+	It("Should return false if string is not RKE2 version", func() {
+		Expect(IsRKE2Version(k8sVersion)).To(BeFalse())
+	})
+})
diff --git a/test/e2e/data/infrastructure/cluster-template-docker.yaml b/test/e2e/data/infrastructure/cluster-template-docker.yaml
index 6fda698..2c8a32a 100644
--- a/test/e2e/data/infrastructure/cluster-template-docker.yaml
+++ b/test/e2e/data/infrastructure/cluster-template-docker.yaml
@@ -99,8 +99,8 @@ metadata:
   name: ${CLUSTER_NAME}-control-plane
 spec: 
   replicas: ${CONTROL_PLANE_MACHINE_COUNT}
+  version: ${KUBERNETES_VERSION}+rke2r1
   agentConfig:
-    version: ${KUBERNETES_VERSION}+rke2r1
     nodeAnnotations:
       test: "true"
   serverConfig:
@@ -135,7 +135,7 @@ spec:
       cluster.x-k8s.io/cluster-name: ${CLUSTER_NAME}
   template:
     spec:
-      version: ${KUBERNETES_VERSION}
+      version: ${KUBERNETES_VERSION}+rke2r1
       clusterName: ${CLUSTER_NAME}
       bootstrap:
         configRef:
@@ -165,6 +165,5 @@ spec:
   template:
     spec:
       agentConfig:
-        version: ${KUBERNETES_VERSION}+rke2r1
         nodeAnnotations:
           test: "true"
diff --git a/test/e2e/helpers.go b/test/e2e/helpers.go
index 28df95b..f48e2f4 100644
--- a/test/e2e/helpers.go
+++ b/test/e2e/helpers.go
@@ -33,14 +33,14 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/klog/v2"
 
+	controlplanev1alpha1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1alpha1"
+	controlplanev1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1beta1"
+	bsutil "github.com/rancher-sandbox/cluster-api-provider-rke2/pkg/util"
+	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 	"sigs.k8s.io/cluster-api/test/framework"
 	"sigs.k8s.io/cluster-api/test/framework/clusterctl"
 	"sigs.k8s.io/cluster-api/util/conditions"
 	"sigs.k8s.io/controller-runtime/pkg/client"
-
-	controlplanev1alpha1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1alpha1"
-	controlplanev1 "github.com/rancher-sandbox/cluster-api-provider-rke2/controlplane/api/v1beta1"
-	clusterv1 "sigs.k8s.io/cluster-api/api/v1beta1"
 )
 
 // NOTE: the code in this file is largely copied from the cluster-api test framework with
@@ -436,7 +436,12 @@ func WaitForClusterToUpgrade(ctx context.Context, input WaitForClusterToUpgradeI
 		}
 
 		for _, machine := range machineList.Items {
-			if machine.Spec.Version != nil && *machine.Spec.Version != input.VersionAfterUpgrade {
+			expectedVersion := input.VersionAfterUpgrade
+			if bsutil.IsRKE2Version(*machine.Spec.Version) {
+				expectedVersion = input.VersionAfterUpgrade + "+rke2r1"
+			}
+
+			if machine.Spec.Version != nil && *machine.Spec.Version != expectedVersion {
 				return false, nil
 			}
 		}