diff --git a/Makefile b/Makefile
index ea086ce5d89..36da33ee4c3 100644
--- a/Makefile
+++ b/Makefile
@@ -44,7 +44,7 @@ BUILD_DATE := $(shell date -u +"%Y-%m-%dT%H:%M:%SZ")
 endif
 LDFLAGS = -X ${PKG}/pkg/smb.driverVersion=${IMAGE_VERSION} -X ${PKG}/pkg/smb.gitCommit=${GIT_COMMIT} -X ${PKG}/pkg/smb.buildDate=${BUILD_DATE}
 EXT_LDFLAGS = -s -w -extldflags "-static"
-E2E_HELM_OPTIONS ?= --set image.smb.repository=$(REGISTRY)/$(IMAGENAME) --set image.smb.tag=$(IMAGE_VERSION)
+E2E_HELM_OPTIONS ?= --set image.smb.repository=$(REGISTRY)/$(IMAGENAME) --set image.smb.tag=$(IMAGE_VERSION) --set controller.runOnControlPlane=true
 E2E_HELM_OPTIONS += ${EXTRA_HELM_OPTIONS}
 # Generate all combination of all OS, ARCH, and OSVERSIONS for iteration
 ALL_OS = linux windows
diff --git a/charts/latest/csi-driver-smb-v0.0.0.tgz b/charts/latest/csi-driver-smb-v0.0.0.tgz
index ccb42500c61..028adf6ea18 100644
Binary files a/charts/latest/csi-driver-smb-v0.0.0.tgz and b/charts/latest/csi-driver-smb-v0.0.0.tgz differ
diff --git a/charts/latest/csi-driver-smb/templates/csi-smb-controller.yaml b/charts/latest/csi-driver-smb/templates/csi-smb-controller.yaml
index 4961f82f1d6..eb33cae9488 100755
--- a/charts/latest/csi-driver-smb/templates/csi-smb-controller.yaml
+++ b/charts/latest/csi-driver-smb/templates/csi-smb-controller.yaml
@@ -24,10 +24,20 @@ spec:
 {{ toYaml .Values.podAnnotations | indent 8 }}
 {{- end }}
     spec:
-{{- with .Values.controller.affinity }}
+    {{- if contains (tpl "{{ .Values.controller.affinity }}" .) "nodeSelectorTerms" }}
+      {{- with .Values.controller.affinity }}
       affinity:
-{{ toYaml . | indent 8 }}
-{{- end }}
+      {{ toYaml . | indent 8 }}
+      {{- end }}
+    {{- else if .Values.controller.runOnControlPlane}}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                - key: node-role.kubernetes.io/control-plane
+                  operator: Exists
+    {{- end }}
       hostNetwork: true
       dnsPolicy: {{ .Values.controller.dnsPolicy }}
       serviceAccountName: {{ .Values.serviceAccount.controller }}
@@ -39,9 +49,6 @@ spec:
         {{- if .Values.controller.runOnMaster}}
         node-role.kubernetes.io/master: ""
         {{- end}}
-        {{- if .Values.controller.runOnControlPlane}}
-        node-role.kubernetes.io/control-plane: ""
-        {{- end}}
       priorityClassName: {{ .Values.priorityClassName | quote }}
       {{- if .Values.securityContext }}
       securityContext: {{- toYaml .Values.securityContext | nindent 8 }}