diff --git a/Makefile b/Makefile
index 10fea9467..b559b7888 100644
--- a/Makefile
+++ b/Makefile
@@ -29,10 +29,9 @@ CSI_IMAGE_TAG ?= $(REGISTRY)/$(IMAGE_NAME):$(IMAGE_VERSION)
 CSI_IMAGE_TAG_LATEST = $(REGISTRY)/$(IMAGE_NAME):latest
 BUILD_DATE ?= $(shell date -u +"%Y-%m-%dT%H:%M:%SZ")
 LDFLAGS ?= "-X ${PKG}/pkg/blob.driverVersion=${IMAGE_VERSION} -X ${PKG}/pkg/blob.gitCommit=${GIT_COMMIT} -X ${PKG}/pkg/blob.buildDate=${BUILD_DATE} -s -w -extldflags '-static'"
+E2E_HELM_OPTIONS ?= --set image.blob.pullPolicy=Always --set image.blob.repository=$(REGISTRY)/$(IMAGE_NAME) --set image.blob.tag=$(IMAGE_VERSION) --set driver.userAgentSuffix="e2e-test" --set node.blobfuseProxy.migrateK8sRepo=true --set controller.runOnControlPlane=true
 ifdef ENABLE_BLOBFUSE_PROXY
-E2E_HELM_OPTIONS ?= --set image.blob.pullPolicy=Always --set image.blob.repository=$(REGISTRY)/$(IMAGE_NAME) --set image.blob.tag=$(IMAGE_VERSION) --set driver.userAgentSuffix="e2e-test" --set controller.logLevel=6 --set node.logLevel=6 --set node.enableBlobfuseProxy=true --set node.blobfuseProxy.migrateK8sRepo=true
-else
-E2E_HELM_OPTIONS ?= --set image.blob.pullPolicy=Always --set image.blob.repository=$(REGISTRY)/$(IMAGE_NAME) --set image.blob.tag=$(IMAGE_VERSION) --set driver.userAgentSuffix="e2e-test" --set node.blobfuseProxy.migrateK8sRepo=true
+E2E_HELM_OPTIONS += --set node.enableBlobfuseProxy=true --set image.blob.pullPolicy=Always --set controller.logLevel=6 --set node.logLevel=6
 endif
 E2E_HELM_OPTIONS += ${EXTRA_HELM_OPTIONS}
 GO111MODULE = on
diff --git a/charts/latest/blob-csi-driver-v1.24.5.tgz b/charts/latest/blob-csi-driver-v1.24.5.tgz
index 1cd32c03f..d926d2ef4 100644
Binary files a/charts/latest/blob-csi-driver-v1.24.5.tgz and b/charts/latest/blob-csi-driver-v1.24.5.tgz differ
diff --git a/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml b/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml
index eefa455fb..6fa49e711 100644
--- a/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml
+++ b/charts/latest/blob-csi-driver/templates/csi-blob-controller.yaml
@@ -28,10 +28,6 @@ spec:
 {{ toYaml .Values.podAnnotations | indent 8 }}
 {{- end }}
     spec:
-{{- with .Values.controller.affinity }}
-      affinity:
-{{ toYaml . | indent 8 }}
-{{- end }}
       {{- if .Values.imagePullSecrets }}
       imagePullSecrets:
 {{ toYaml .Values.imagePullSecrets | indent 8 }}
@@ -40,12 +36,27 @@ spec:
       serviceAccountName: {{ .Values.serviceAccount.controller }}
       nodeSelector:
         kubernetes.io/os: linux
-        {{- if .Values.controller.runOnMaster}}
-        node-role.kubernetes.io/master: ""
-        {{- end}}
-        {{- if .Values.controller.runOnControlPlane}}
-        node-role.kubernetes.io/control-plane: ""
-        {{- end}}
+    # runOnControlPlane=true or runOnMaster=true only takes effect if affinity is not set
+    {{- if contains (tpl "{{ .Values.controller.affinity }}" .) "nodeSelectorTerms" }}
+      {{- with .Values.controller.affinity }}
+      affinity:
+      {{ toYaml . | indent 8 }}
+      {{- end }}
+    {{- else if or .Values.controller.runOnControlPlane .Values.controller.runOnMaster}}
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+              - matchExpressions:
+                {{- if .Values.controller.runOnControlPlane}}
+                - key: node-role.kubernetes.io/control-plane
+                  operator: Exists
+                {{- end}}
+                {{- if .Values.controller.runOnMaster}}
+                - key: node-role.kubernetes.io/master
+                  operator: Exists
+                {{- end}}
+    {{- end }}
 {{- with .Values.controller.nodeSelector }}
 {{ toYaml . | indent 8 }}
 {{- end }}