Repository: kubernetes-sigs/azuredisk-csi-driver
Pull request #727 (2021-02-09 19:10:51)

chore: Add side-by-side implementation of v2 driver to ease integration with ongoing v1 work
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://github.com/kubernetes/community/blob/master/contributors/devel/sig-release/release.md#issue-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://github.com/kubernetes/community/blob/master/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. Follow the instructions for writing a release note: https://git.k8s.io/community/contributors/guide/release-notes.md
6. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

**What type of PR is this?**

/kind cleanup

<!--
Add one of the following kinds:
/kind bug
/kind test
/kind cleanup
/kind documentation
/kind feature
/kind design

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

**What this PR does / why we need it**:

This PR adds a side-by-side implementation of the v2 driver that is conditionally compiled (through the "azurediskv2" tag) in order to ease integration with ongoing v1 work with v2 work.

**Which issue(s) this PR fixes**:
<!-- 
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #

**Requirements**:
- [X] uses [conventional commit messages](https://www.conventionalcommits.org/)
  <!-- Common commit types:
        build: Build 🏭
        chore: Maintenance 🔧
        ci: Continuous Integration 💜
        docs: Documentation 📘
        feat: Features 🌈
        fix: Bug Fixes 🐞
        perf: Performance Improvements 🚀
        refactor: Code Refactoring 💎
        revert: Revert Change ◀️
        style: Code Style 🎶
        security: Security Fix 🛡️
        test: Testing 💚 -->
- [ ] includes documentation
- [X] adds unit tests
- [ ] tested upgrade from previous version

**Special notes for your reviewer**:

@andyzhangx This is the PR we discussed offline for staging v2 development. The v2 implementation is conditionally compiled through the "azurediskv2" tag. There is also a "--temp-use-driver-v2" command-line switch to enable side-by-side testing in the same binary. In this PR, the v1 and v2 implementations are the same (xxx_v2.go files are copies of their original counterparts). I would appreciate your feedback on this approach and your recommendations for integrating v2 testing into the integration and e2e testing pipelines.

**Release note**:
```
none
```


k8s-ci-robot: Hi @edreed. Thanks for your PR.

I'm waiting for a [kubernetes-sigs](https://github.com/orgs/kubernetes-sigs/people) member to verify that this patch is reasonable to test. If it is, they should reply with `/ok-to-test` on its own line. Until that is done, I will not automatically test new commits in this PR, but the usual testing commands by org members will still work. Regular contributors should [join the org](https://git.k8s.io/community/community-membership.md#member) to skip this step.

Once the patch is verified, the new status will be reflected by the `ok-to-test` label.

I understand the commands that are listed [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fazuredisk-csi-driver).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>

k8s-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/kubernetes-sigs/azuredisk-csi-driver/pull/727#" title="Author self-approved">edreed</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **feiskyer** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @feiskyer` in a comment when ready.

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=kubernetes-sigs%2Fazuredisk-csi-driver).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/kubernetes-sigs/azuredisk-csi-driver/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["feiskyer"]} -->
k8s-ci-robot: @edreed: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-azuredisk-csi-driver-e2e-windows | 062bc6bf77afdb0426c3981b04141b6c4c032ab6 | [link](https://prow.k8s.io/view/gs/kubernetes-jenkins/pr-logs/pull/kubernetes-sigs_azuredisk-csi-driver/727/pull-azuredisk-csi-driver-e2e-windows/1359313001865285632) | `/test pull-azuredisk-csi-driver-e2e-windows`

[Full PR test history](https://prow.k8s.io/pr-history?org=kubernetes-sigs&repo=azuredisk-csi-driver&pr=727). [Your PR dashboard](https://prow.k8s.io/pr?query=is%3Apr%20state%3Aopen%20author%3Aedreed). Please help us cut down on flakes by [linking to](https://git.k8s.io/community/contributors/devel/sig-testing/flaky-tests.md#filing-issues-for-flaky-tests) an [open issue](https://github.com/kubernetes-sigs/azuredisk-csi-driver/issues?q=is:issue+is:open) when you hit one in your PR.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
edreed: I rebased to pick up the fix for the Windows unit tests, and all unit tests are now passing.
k8s-ci-robot: @edreed: The following tests **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-azuredisk-csi-driver-e2e-windows | 062bc6bf77afdb0426c3981b04141b6c4c032ab6 | [link](https://prow.k8s.io/view/gs/kubernetes-jenkins/pr-logs/pull/kubernetes-sigs_azuredisk-csi-driver/727/pull-azuredisk-csi-driver-e2e-windows/1359313001865285632) | `/test pull-azuredisk-csi-driver-e2e-windows`
pull-azuredisk-csi-driver-e2e-single-az | 184fef0b15c0fa14e51ed45fab9e146289506acf | [link](https://prow.k8s.io/view/gs/kubernetes-jenkins/pr-logs/pull/kubernetes-sigs_azuredisk-csi-driver/727/pull-azuredisk-csi-driver-e2e-single-az/1362979295990910976) | `/test pull-azuredisk-csi-driver-e2e-single-az`

[Full PR test history](https://prow.k8s.io/pr-history?org=kubernetes-sigs&repo=azuredisk-csi-driver&pr=727). [Your PR dashboard](https://prow.k8s.io/pr?query=is%3Apr%20state%3Aopen%20author%3Aedreed). Please help us cut down on flakes by [linking to](https://git.k8s.io/community/contributors/devel/sig-testing/flaky-tests.md#filing-issues-for-flaky-tests) an [open issue](https://github.com/kubernetes-sigs/azuredisk-csi-driver/issues?q=is:issue+is:open) when you hit one in your PR.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
edreed: /retest
k8s-ci-robot: @edreed: PR needs rebase.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
k8s-ci-robot: @edreed: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-azuredisk-csi-driver-sanity-v2 | 80b3562e2315df1281477d4d4e618096a833a3f0 | [link](https://prow.k8s.io/view/gs/kubernetes-jenkins/pr-logs/pull/kubernetes-sigs_azuredisk-csi-driver/727/pull-azuredisk-csi-driver-sanity-v2/1364380931917352960) | `/test pull-azuredisk-csi-driver-sanity-v2`

[Full PR test history](https://prow.k8s.io/pr-history?org=kubernetes-sigs&repo=azuredisk-csi-driver&pr=727). [Your PR dashboard](https://prow.k8s.io/pr?query=is%3Apr%20state%3Aopen%20author%3Aedreed). Please help us cut down on flakes by [linking to](https://git.k8s.io/community/contributors/devel/sig-testing/flaky-tests.md#filing-issues-for-flaky-tests) an [open issue](https://github.com/kubernetes-sigs/azuredisk-csi-driver/issues?q=is:issue+is:open) when you hit one in your PR.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
andyzhangx: /test pull-azuredisk-csi-driver-sanity-v2
k8s-ci-robot: @edreed: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-azuredisk-csi-driver-sanity-v2 | 80b3562e2315df1281477d4d4e618096a833a3f0 | [link](https://prow.k8s.io/view/gs/kubernetes-jenkins/pr-logs/pull/kubernetes-sigs_azuredisk-csi-driver/727/pull-azuredisk-csi-driver-sanity-v2/1364389647974666240) | `/test pull-azuredisk-csi-driver-sanity-v2`

[Full PR test history](https://prow.k8s.io/pr-history?org=kubernetes-sigs&repo=azuredisk-csi-driver&pr=727). [Your PR dashboard](https://prow.k8s.io/pr?query=is%3Apr%20state%3Aopen%20author%3Aedreed). Please help us cut down on flakes by [linking to](https://git.k8s.io/community/contributors/devel/sig-testing/flaky-tests.md#filing-issues-for-flaky-tests) an [open issue](https://github.com/kubernetes-sigs/azuredisk-csi-driver/issues?q=is:issue+is:open) when you hit one in your PR.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
edreed: /test pull-azuredisk-csi-driver-sanity-v2

Patch:
diff --git a/Makefile b/Makefile
index 399998e88..a0c1e8c22 100644
--- a/Makefile
+++ b/Makefile
@@ -72,17 +72,32 @@ verify: unit-test
 	go vet ./pkg/...
 
 .PHONY: unit-test
-unit-test:
+unit-test: unit-test-v1 unit-test-v2
+
+.PHONY: unit-test-v1
+unit-test-v1:
 	go test -v -cover ./pkg/... ./test/utils/credentials
 
+.PHONY: unit-test-v2
+unit-test-v2:
+	go test -v -cover -tags azurediskv2 ./pkg/azuredisk --temp-use-driver-v2
+
 .PHONY: sanity-test
 sanity-test: azuredisk
 	go test -v -timeout=30m ./test/sanity
 
+.PHONY: sanity-test-v2
+sanity-test-v2: azuredisk-v2
+	go test -v -timeout=30m ./test/sanity --temp-use-driver-v2
+
 .PHONY: integration-test
 integration-test: azuredisk
 	go test -v -timeout=30m ./test/integration
 
+.PHONY: integration-test-v2
+integration-test-v2: azuredisk-v2
+	go test -v -timeout=30m ./test/integration --temp-use-driver-v2
+
 .PHONY: e2e-test
 e2e-test:
 	go test -v -timeout=0 ./test/e2e ${GINKGO_FLAGS}
@@ -119,6 +134,10 @@ e2e-teardown:
 azuredisk:
 	CGO_ENABLED=0 GOOS=linux go build -a -ldflags ${LDFLAGS} -mod vendor -o _output/azurediskplugin ./pkg/azurediskplugin
 
+.PHONY: azuredisk-v2
+azuredisk-v2:
+	CGO_ENABLED=0 GOOS=linux go build -a -ldflags ${LDFLAGS} -mod vendor -tags azurediskv2 -o _output/azurediskpluginv2 ./pkg/azurediskplugin
+
 .PHONY: azuredisk-windows
 azuredisk-windows:
 	CGO_ENABLED=0 GOOS=windows go build -a -ldflags ${LDFLAGS} -mod vendor -o _output/azurediskplugin.exe ./pkg/azurediskplugin
diff --git a/pkg/azuredisk/azuredisk.go b/pkg/azuredisk/azuredisk.go
index fe00271ff..4b3f654ab 100644
--- a/pkg/azuredisk/azuredisk.go
+++ b/pkg/azuredisk/azuredisk.go
@@ -38,6 +38,7 @@ import (
 	csicommon "sigs.k8s.io/azuredisk-csi-driver/pkg/csi-common"
 	"sigs.k8s.io/azuredisk-csi-driver/pkg/mounter"
 	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+	"sigs.k8s.io/cloud-provider-azure/pkg/provider"
 	azure "sigs.k8s.io/cloud-provider-azure/pkg/provider"
 )
 
@@ -92,18 +93,32 @@ var (
 	diskURISupportedManaged = []string{"/subscriptions/{sub-id}/resourcegroups/{group-name}/providers/microsoft.compute/disks/{disk-id}"}
 )
 
-// Driver implements all interfaces of CSI drivers
-type Driver struct {
+// CSIDriver defines the interface for a CSI driver.
+type CSIDriver interface {
+	csi.ControllerServer
+	csi.NodeServer
+	csi.IdentityServer
+
+	Run(endpoint, kubeconfig string, testMode bool)
+}
+
+// DriverCore contains fields common to both the V1 and V2 driver, and implements all interfaces of CSI drivers
+type DriverCore struct {
 	csicommon.CSIDriver
-	cloud             *azure.Cloud
-	supportAzureStack bool
-	mounter           *mount.SafeFormatAndMount
+	cloud   *azure.Cloud
+	mounter *mount.SafeFormatAndMount
+}
+
+// Driver is the v1 implementation of the Azure Disk CSI Driver.
+type Driver struct {
+	DriverCore
 	volumeLocks       *volumehelper.VolumeLocks
+	supportAzureStack bool
 }
 
-// NewDriver Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
+// newDriverV1 Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
 // does not support optional driver plugin info manifest field. Refer to CSI spec for more details.
-func NewDriver(nodeID string, supportAzureStack bool) *Driver {
+func newDriverV1(nodeID string, supportAzureStack bool) *Driver {
 	driver := Driver{}
 	driver.Name = DriverName
 	driver.Version = driverVersion
@@ -325,3 +340,48 @@ func IsCorruptedDir(dir string) bool {
 	fmt.Printf("IsCorruptedDir(%s) returned with error: %v", dir, pathErr)
 	return pathErr != nil && mount.IsCorruptedMnt(pathErr)
 }
+
+// setControllerCapabilities sets the controller capabilities field. It is intended for use with unit tests.
+func (d *DriverCore) setControllerCapabilities(caps []*csi.ControllerServiceCapability) {
+	d.Cap = caps
+}
+
+// setNodeCapabilities sets the node capabilities field. It is intended for use with unit tests.
+func (d *DriverCore) setNodeCapabilities(nodeCaps []*csi.NodeServiceCapability) {
+	d.NSCap = nodeCaps
+}
+
+// setName sets the Name field. It is intended for use with unit tests.
+func (d *DriverCore) setName(name string) {
+	d.Name = name
+}
+
+// setName sets the NodeId field. It is intended for use with unit tests.
+func (d *DriverCore) setNodeID(nodeID string) {
+	d.NodeID = nodeID
+}
+
+// setName sets the Version field. It is intended for use with unit tests.
+func (d *DriverCore) setVersion(version string) {
+	d.Version = version
+}
+
+// getCloud returns the value of the cloud field. It is intended for use with unit tests.
+func (d *DriverCore) getCloud() *provider.Cloud {
+	return d.cloud
+}
+
+// setCloud sets the cloud field. It is intended for use with unit tests.
+func (d *DriverCore) setCloud(cloud *provider.Cloud) {
+	d.cloud = cloud
+}
+
+// getMounter returns the value of the mounter field. It is intended for use with unit tests.
+func (d *DriverCore) getMounter() *mount.SafeFormatAndMount {
+	return d.mounter
+}
+
+// setMounter sets the mounter field. It is intended for use with unit tests.
+func (d *DriverCore) setMounter(mounter *mount.SafeFormatAndMount) {
+	d.mounter = mounter
+}
diff --git a/pkg/azuredisk/azuredisk_v1.go b/pkg/azuredisk/azuredisk_v1.go
new file mode 100644
index 000000000..a7fe1954f
--- /dev/null
+++ b/pkg/azuredisk/azuredisk_v1.go
@@ -0,0 +1,25 @@
+// +build !azurediskv2
+
+/*
+Copyright 2017 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+// NewDriver Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
+// does not support optional driver plugin info manifest field. Refer to CSI spec for more details.
+func NewDriver(nodeID string, supportAzureStack bool) CSIDriver {
+	return newDriverV1(nodeID, supportAzureStack)
+}
diff --git a/pkg/azuredisk/azuredisk_v2.go b/pkg/azuredisk/azuredisk_v2.go
new file mode 100644
index 000000000..7fe5e96d7
--- /dev/null
+++ b/pkg/azuredisk/azuredisk_v2.go
@@ -0,0 +1,151 @@
+// +build azurediskv2
+
+/*
+Copyright 2017 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import (
+	"context"
+	"flag"
+	"reflect"
+
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-30/compute"
+	"github.com/container-storage-interface/spec/lib/go/csi"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"k8s.io/klog/v2"
+
+	csicommon "sigs.k8s.io/azuredisk-csi-driver/pkg/csi-common"
+	"sigs.k8s.io/azuredisk-csi-driver/pkg/mounter"
+	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+)
+
+var useDriverV2 = flag.Bool("temp-use-driver-v2", false, "A temporary flag to enable early test and development of Azure Disk CSI Driver V2. This will be removed in the future.")
+
+// DriverV2 implements all interfaces of CSI drivers
+type DriverV2 struct {
+	DriverCore
+	volumeLocks       *volumehelper.VolumeLocks
+	supportAzureStack bool
+}
+
+// NewDriver creates a Driver or DriverV2 object depending on the --temp-use-driver-v2 flag.
+func NewDriver(nodeID string, supportAzureStack bool) CSIDriver {
+	var d CSIDriver
+
+	if !*useDriverV2 {
+		d = newDriverV1(nodeID, supportAzureStack)
+	} else {
+		d = newDriverV2(nodeID, supportAzureStack)
+	}
+
+	return d
+}
+
+// newDriverV2 Creates a NewCSIDriver object. Assumes vendor version is equal to driver version &
+// does not support optional driver plugin info manifest field. Refer to CSI spec for more details.
+func newDriverV2(nodeID string, supportAzureStack bool) *DriverV2 {
+	klog.Warning("Using DriverV2")
+	driver := DriverV2{}
+	driver.Name = DriverName
+	driver.Version = driverVersion
+	driver.NodeID = nodeID
+	driver.volumeLocks = volumehelper.NewVolumeLocks()
+	driver.supportAzureStack = supportAzureStack
+	return &driver
+}
+
+// Run driver initialization
+func (d *DriverV2) Run(endpoint, kubeconfig string, testBool bool) {
+	versionMeta, err := GetVersionYAML()
+	if err != nil {
+		klog.Fatalf("%v", err)
+	}
+	klog.Infof("\nDRIVER INFORMATION:\n-------------------\n%s\n\nStreaming logs below:", versionMeta)
+	cloud, err := GetCloudProvider(kubeconfig)
+	if err != nil || cloud.TenantID == "" || cloud.SubscriptionID == "" {
+		klog.Fatalf("failed to get Azure Cloud Provider, error: %v", err)
+	}
+	d.cloud = cloud
+
+	if d.NodeID == "" {
+		// Disable UseInstanceMetadata for controller to mitigate a timeout issue using IMDS
+		// https://github.com/kubernetes-sigs/azuredisk-csi-driver/issues/168
+		klog.Infoln("disable UseInstanceMetadata for controller")
+		d.cloud.Config.UseInstanceMetadata = false
+	}
+
+	d.mounter, err = mounter.NewSafeMounter()
+	if err != nil {
+		klog.Fatalf("Failed to get safe mounter. Error: %v", err)
+	}
+
+	d.AddControllerServiceCapabilities(
+		[]csi.ControllerServiceCapability_RPC_Type{
+			csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME,
+			csi.ControllerServiceCapability_RPC_PUBLISH_UNPUBLISH_VOLUME,
+			csi.ControllerServiceCapability_RPC_CREATE_DELETE_SNAPSHOT,
+			csi.ControllerServiceCapability_RPC_LIST_SNAPSHOTS,
+			csi.ControllerServiceCapability_RPC_CLONE_VOLUME,
+			csi.ControllerServiceCapability_RPC_EXPAND_VOLUME,
+			csi.ControllerServiceCapability_RPC_LIST_VOLUMES,
+			csi.ControllerServiceCapability_RPC_LIST_VOLUMES_PUBLISHED_NODES,
+		})
+	d.AddVolumeCapabilityAccessModes([]csi.VolumeCapability_AccessMode_Mode{csi.VolumeCapability_AccessMode_SINGLE_NODE_WRITER})
+	d.AddNodeServiceCapabilities([]csi.NodeServiceCapability_RPC_Type{
+		csi.NodeServiceCapability_RPC_STAGE_UNSTAGE_VOLUME,
+		csi.NodeServiceCapability_RPC_EXPAND_VOLUME,
+		csi.NodeServiceCapability_RPC_GET_VOLUME_STATS,
+	})
+
+	s := csicommon.NewNonBlockingGRPCServer()
+	// Driver d act as IdentityServer, ControllerServer and NodeServer
+	s.Start(endpoint, d, d, d, testBool)
+	s.Wait()
+}
+
+func (d *DriverV2) checkDiskExists(ctx context.Context, diskURI string) error {
+	diskName, err := GetDiskName(diskURI)
+	if err != nil {
+		return err
+	}
+
+	resourceGroup, err := GetResourceGroupFromURI(diskURI)
+	if err != nil {
+		return err
+	}
+
+	if _, rerr := d.cloud.DisksClient.Get(ctx, resourceGroup, diskName); rerr != nil {
+		return rerr.Error()
+	}
+
+	return nil
+}
+
+func (d *DriverV2) checkDiskCapacity(ctx context.Context, resourceGroup, diskName string, requestGiB int) (bool, error) {
+	disk, err := d.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
+	// Because we can not judge the reason of the error. Maybe the disk does not exist.
+	// So here we do not handle the error.
+	if err == nil {
+		if !reflect.DeepEqual(disk, compute.Disk{}) && disk.DiskSizeGB != nil && int(*disk.DiskSizeGB) != requestGiB {
+			return false, status.Errorf(codes.AlreadyExists, "the request volume already exists, but its capacity(%v) is different from (%v)", *disk.DiskProperties.DiskSizeGB, requestGiB)
+		}
+	}
+	return true, nil
+}
diff --git a/pkg/azuredisk/controllerserver_v2.go b/pkg/azuredisk/controllerserver_v2.go
new file mode 100644
index 000000000..d3583c435
--- /dev/null
+++ b/pkg/azuredisk/controllerserver_v2.go
@@ -0,0 +1,953 @@
+// +build azurediskv2
+
+/*
+Copyright 2017 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import (
+	"context"
+	"fmt"
+	"path"
+	"sort"
+	"strconv"
+	"strings"
+
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2020-06-30/compute"
+	"github.com/container-storage-interface/spec/lib/go/csi"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	cloudprovider "k8s.io/cloud-provider"
+	volerr "k8s.io/cloud-provider/volume/errors"
+	"k8s.io/klog/v2"
+
+	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+	"sigs.k8s.io/cloud-provider-azure/pkg/metrics"
+	azure "sigs.k8s.io/cloud-provider-azure/pkg/provider"
+)
+
+// CreateVolume provisions an azure disk
+func (d *DriverV2) CreateVolume(ctx context.Context, req *csi.CreateVolumeRequest) (*csi.CreateVolumeResponse, error) {
+	if err := d.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
+		klog.Errorf("invalid create volume req: %v", req)
+		return nil, err
+	}
+
+	name := req.GetName()
+	if len(name) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "CreateVolume Name must be provided")
+	}
+	volCaps := req.GetVolumeCapabilities()
+	if len(volCaps) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "CreateVolume Volume capabilities must be provided")
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(name); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, name)
+	}
+	defer d.volumeLocks.Release(name)
+
+	capacityBytes := req.GetCapacityRange().GetRequiredBytes()
+	volSizeBytes := int64(capacityBytes)
+	requestGiB := int(volumehelper.RoundUpGiB(volSizeBytes))
+	if requestGiB < minimumDiskSizeGiB {
+		requestGiB = minimumDiskSizeGiB
+	}
+
+	maxVolSize := int(volumehelper.RoundUpGiB(req.GetCapacityRange().GetLimitBytes()))
+	if (maxVolSize > 0) && (maxVolSize < requestGiB) {
+		return nil, status.Error(codes.InvalidArgument, "After round-up, volume size exceeds the limit specified")
+	}
+
+	var (
+		location                string
+		storageAccountType      string
+		cachingMode             v1.AzureDataDiskCachingMode
+		err                     error
+		resourceGroup           string
+		diskIopsReadWrite       string
+		diskMbpsReadWrite       string
+		logicalSectorSize       int
+		diskName                string
+		diskEncryptionSetID     string
+		customTags              string
+		writeAcceleratorEnabled string
+		maxShares               int
+	)
+
+	parameters := req.GetParameters()
+	if parameters == nil {
+		parameters = make(map[string]string)
+	}
+	for k, v := range parameters {
+		switch strings.ToLower(k) {
+		case skuNameField:
+			storageAccountType = v
+		case locationField:
+			location = v
+		case storageAccountTypeField:
+			storageAccountType = v
+		case cachingModeField:
+			cachingMode = v1.AzureDataDiskCachingMode(v)
+		case resourceGroupField:
+			resourceGroup = v
+		case diskIOPSReadWriteField:
+			diskIopsReadWrite = v
+		case diskMBPSReadWriteField:
+			diskMbpsReadWrite = v
+		case logicalSectorSizeField:
+			logicalSectorSize, err = strconv.Atoi(v)
+			if err != nil {
+				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("parse %s failed with error: %v", v, err))
+			}
+		case diskNameField:
+			diskName = v
+		case desIDField:
+			diskEncryptionSetID = v
+		case tagsField:
+			customTags = v
+		case azure.WriteAcceleratorEnabled:
+			writeAcceleratorEnabled = v
+		case maxSharesField:
+			maxShares, err = strconv.Atoi(v)
+			if err != nil {
+				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("parse %s failed with error: %v", v, err))
+			}
+			if maxShares < 1 {
+				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("parse %s returned with invalid value: %d", v, maxShares))
+			}
+		default:
+			//don't return error here since there are some parameters(e.g. fsType) used in disk mount process
+			//return nil, fmt.Errorf("AzureDisk - invalid option %s in storage class", k)
+		}
+	}
+
+	if IsAzureStackCloud(d.cloud.Config.Cloud, d.supportAzureStack) {
+		if maxShares > 1 {
+			return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("Invalid maxShares value: %d as Azure Stack does not support shared disk.", maxShares))
+		}
+	}
+
+	if maxShares < 2 {
+		for _, c := range volCaps {
+			mode := c.GetAccessMode().Mode
+			if mode != csi.VolumeCapability_AccessMode_SINGLE_NODE_WRITER &&
+				mode != csi.VolumeCapability_AccessMode_SINGLE_NODE_READER_ONLY {
+				return nil, status.Error(codes.InvalidArgument, fmt.Sprintf("Volume capability(%v) not supported", mode))
+			}
+		}
+	}
+
+	if diskName == "" {
+		diskName = name
+	}
+	diskName = getValidDiskName(diskName)
+
+	if resourceGroup == "" {
+		resourceGroup = d.cloud.ResourceGroup
+	}
+
+	// normalize values
+	skuName, err := normalizeStorageAccountType(storageAccountType, d.cloud.Config.Cloud, d.supportAzureStack)
+	if err != nil {
+		return nil, err
+	}
+
+	if _, err = normalizeCachingMode(cachingMode); err != nil {
+		return nil, err
+	}
+
+	selectedAvailabilityZone := pickAvailabilityZone(req.GetAccessibilityRequirements(), d.cloud.Location)
+
+	if ok, err := d.checkDiskCapacity(ctx, resourceGroup, diskName, requestGiB); !ok {
+		return nil, err
+	}
+
+	customTagsMap, err := volumehelper.ConvertTagsToMap(customTags)
+	if err != nil {
+		return nil, err
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_create_volume", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("begin to create azure disk(%s) account type(%s) rg(%s) location(%s) size(%d) selectedAvailabilityZone(%v) maxShares(%d)",
+		diskName, skuName, resourceGroup, location, requestGiB, selectedAvailabilityZone, maxShares)
+
+	tags := make(map[string]string)
+	contentSource := &csi.VolumeContentSource{}
+	for k, v := range customTagsMap {
+		tags[k] = v
+	}
+	tags[CreatedForPVNameKey] = name
+
+	if strings.EqualFold(writeAcceleratorEnabled, "true") {
+		tags[azure.WriteAcceleratorEnabled] = "true"
+	}
+	sourceID := ""
+	sourceType := ""
+	content := req.GetVolumeContentSource()
+	if content != nil {
+		if content.GetSnapshot() != nil {
+			sourceID = content.GetSnapshot().GetSnapshotId()
+			sourceType = sourceSnapshot
+			contentSource = &csi.VolumeContentSource{
+				Type: &csi.VolumeContentSource_Snapshot{
+					Snapshot: &csi.VolumeContentSource_SnapshotSource{
+						SnapshotId: sourceID,
+					},
+				},
+			}
+		} else {
+			sourceID = content.GetVolume().GetVolumeId()
+			sourceType = sourceVolume
+			contentSource = &csi.VolumeContentSource{
+				Type: &csi.VolumeContentSource_Volume{
+					Volume: &csi.VolumeContentSource_VolumeSource{
+						VolumeId: sourceID,
+					},
+				},
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			if sourceGiB, _ := d.GetSourceDiskSize(ctx, resourceGroup, path.Base(sourceID), 0, sourceDiskSearchMaxDepth); sourceGiB != nil && *sourceGiB < int32(requestGiB) {
+				parameters[resizeRequired] = strconv.FormatBool(true)
+			}
+			cancel()
+		}
+	}
+
+	volumeOptions := &azure.ManagedDiskOptions{
+		DiskName:            diskName,
+		StorageAccountType:  skuName,
+		ResourceGroup:       resourceGroup,
+		PVCName:             "",
+		SizeGB:              requestGiB,
+		Tags:                tags,
+		AvailabilityZone:    selectedAvailabilityZone,
+		DiskIOPSReadWrite:   diskIopsReadWrite,
+		DiskMBpsReadWrite:   diskMbpsReadWrite,
+		SourceResourceID:    sourceID,
+		SourceType:          sourceType,
+		DiskEncryptionSetID: diskEncryptionSetID,
+		MaxShares:           int32(maxShares),
+		LogicalSectorSize:   int32(logicalSectorSize),
+	}
+	diskURI, err := d.cloud.CreateManagedDisk(volumeOptions)
+	if err != nil {
+		if strings.Contains(err.Error(), NotFound) {
+			return nil, status.Error(codes.NotFound, err.Error())
+		}
+		return nil, err
+	}
+
+	isOperationSucceeded = true
+	klog.V(2).Infof("create azure disk(%s) account type(%s) rg(%s) location(%s) size(%d) tags(%s) successfully", diskName, skuName, resourceGroup, location, requestGiB, tags)
+
+	return &csi.CreateVolumeResponse{
+		Volume: &csi.Volume{
+			VolumeId:      diskURI,
+			CapacityBytes: capacityBytes,
+			VolumeContext: parameters,
+			ContentSource: contentSource,
+			AccessibleTopology: []*csi.Topology{
+				{
+					Segments: map[string]string{topologyKey: selectedAvailabilityZone},
+				},
+			},
+		},
+	}, nil
+}
+
+// DeleteVolume delete an azure disk
+func (d *DriverV2) DeleteVolume(ctx context.Context, req *csi.DeleteVolumeRequest) (*csi.DeleteVolumeResponse, error) {
+	volumeID := req.GetVolumeId()
+	if len(volumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
+	}
+
+	if err := d.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME); err != nil {
+		return nil, fmt.Errorf("invalid delete volume req: %v", req)
+	}
+	diskURI := volumeID
+
+	if err := isValidDiskURI(diskURI); err != nil {
+		klog.Errorf("validateDiskURI(%s) in DeleteVolume failed with error: %v", diskURI, err)
+		return &csi.DeleteVolumeResponse{}, nil
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(volumeID); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, volumeID)
+	}
+	defer d.volumeLocks.Release(volumeID)
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_delete_volume", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("deleting azure disk(%s)", diskURI)
+	err := d.cloud.DeleteManagedDisk(diskURI)
+	klog.V(2).Infof("delete azure disk(%s) returned with %v", diskURI, err)
+	isOperationSucceeded = (err == nil)
+	return &csi.DeleteVolumeResponse{}, err
+}
+
+// ControllerGetVolume get volume
+func (d *DriverV2) ControllerGetVolume(context.Context, *csi.ControllerGetVolumeRequest) (*csi.ControllerGetVolumeResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "")
+}
+
+// ControllerPublishVolume attach an azure disk to a required node
+func (d *DriverV2) ControllerPublishVolume(ctx context.Context, req *csi.ControllerPublishVolumeRequest) (*csi.ControllerPublishVolumeResponse, error) {
+	diskURI := req.GetVolumeId()
+	if len(diskURI) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID not provided")
+	}
+
+	volCap := req.GetVolumeCapability()
+	if volCap == nil {
+		return nil, status.Error(codes.InvalidArgument, "Volume capability not provided")
+	}
+
+	caps := []*csi.VolumeCapability{volCap}
+	if !isValidVolumeCapabilities(caps) {
+		return nil, status.Error(codes.InvalidArgument, "Volume capability not supported")
+	}
+
+	err := d.checkDiskExists(ctx, diskURI)
+	if err != nil {
+		return nil, status.Error(codes.NotFound, fmt.Sprintf("Volume not found, failed with error: %v", err))
+	}
+
+	nodeID := req.GetNodeId()
+	if len(nodeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Node ID not provided")
+	}
+
+	nodeName := types.NodeName(nodeID)
+	diskName, err := GetDiskName(diskURI)
+	if err != nil {
+		return nil, err
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_publish_volume", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("GetDiskLun returned: %v. Initiating attaching volume %q to node %q.", err, diskURI, nodeName)
+
+	lun, err := d.cloud.GetDiskLun(diskName, diskURI, nodeName)
+	if err == cloudprovider.InstanceNotFound {
+		return nil, status.Error(codes.NotFound, fmt.Sprintf("failed to get azure instance id for node %q (%v)", nodeName, err))
+	}
+
+	if err == nil {
+		// Volume is already attached to node.
+		klog.V(2).Infof("Attach operation is successful. volume %q is already attached to node %q at lun %d.", diskURI, nodeName, lun)
+	} else {
+		var cachingMode compute.CachingTypes
+		if cachingMode, err = getCachingMode(req.GetVolumeContext()); err != nil {
+			return nil, err
+		}
+		klog.V(2).Infof("Trying to attach volume %q to node %q", diskURI, nodeName)
+
+		lun, err = d.cloud.AttachDisk(true, diskName, diskURI, nodeName, cachingMode)
+		if err == nil {
+			klog.V(2).Infof("Attach operation successful: volume %q attached to node %q.", diskURI, nodeName)
+		} else {
+			if derr, ok := err.(*volerr.DanglingAttachError); ok {
+				klog.Warningf("volume %q is already attached to node %q, try detach first", diskURI, derr.CurrentNode)
+				if err = d.cloud.DetachDisk(diskName, diskURI, derr.CurrentNode); err != nil {
+					return nil, status.Errorf(codes.Internal, "Could not detach volume %q from node %q: %v", diskURI, derr.CurrentNode, err)
+				}
+				klog.V(2).Infof("Trying to attach volume %q to node %q again", diskURI, nodeName)
+				lun, err = d.cloud.AttachDisk(true, diskName, diskURI, nodeName, cachingMode)
+			}
+			if err != nil {
+				klog.Errorf("Attach volume %q to instance %q failed with %v", diskURI, nodeName, err)
+				return nil, fmt.Errorf("Attach volume %q to instance %q failed with %v", diskURI, nodeName, err)
+			}
+		}
+		klog.V(2).Infof("attach volume %q to node %q successfully", diskURI, nodeName)
+	}
+
+	pvInfo := map[string]string{LUN: strconv.Itoa(int(lun))}
+	isOperationSucceeded = true
+	return &csi.ControllerPublishVolumeResponse{PublishContext: pvInfo}, nil
+}
+
+// ControllerUnpublishVolume detach an azure disk from a required node
+func (d *DriverV2) ControllerUnpublishVolume(ctx context.Context, req *csi.ControllerUnpublishVolumeRequest) (*csi.ControllerUnpublishVolumeResponse, error) {
+	diskURI := req.GetVolumeId()
+	if len(diskURI) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID not provided")
+	}
+
+	nodeID := req.GetNodeId()
+	if len(nodeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Node ID not provided")
+	}
+	nodeName := types.NodeName(nodeID)
+
+	diskName, err := GetDiskName(diskURI)
+	if err != nil {
+		return nil, err
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_unpublish_volume", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("Trying to detach volume %s from node %s", diskURI, nodeID)
+
+	if err := d.cloud.DetachDisk(diskName, diskURI, nodeName); err != nil {
+		if strings.Contains(err.Error(), errDiskNotFound) {
+			klog.Warningf("volume %s already detached from node %s", diskURI, nodeID)
+		} else {
+			return nil, status.Errorf(codes.Internal, "Could not detach volume %q from node %q: %v", diskURI, nodeID, err)
+		}
+	}
+	klog.V(2).Infof("detach volume %s from node %s successfully", diskURI, nodeID)
+	isOperationSucceeded = true
+
+	return &csi.ControllerUnpublishVolumeResponse{}, nil
+}
+
+// ValidateVolumeCapabilities return the capabilities of the volume
+func (d *DriverV2) ValidateVolumeCapabilities(ctx context.Context, req *csi.ValidateVolumeCapabilitiesRequest) (*csi.ValidateVolumeCapabilitiesResponse, error) {
+	diskURI := req.GetVolumeId()
+	if len(diskURI) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
+	}
+	if req.GetVolumeCapabilities() == nil {
+		return nil, status.Error(codes.InvalidArgument, "Volume capabilities missing in request")
+	}
+
+	err := d.checkDiskExists(ctx, diskURI)
+	if err != nil {
+		return nil, status.Error(codes.NotFound, fmt.Sprintf("Volume not found, failed with error: %v", err))
+	}
+
+	for _, cap := range req.VolumeCapabilities {
+		if cap.GetAccessMode().GetMode() != csi.VolumeCapability_AccessMode_SINGLE_NODE_WRITER {
+			return &csi.ValidateVolumeCapabilitiesResponse{Message: ""}, nil
+		}
+	}
+	return &csi.ValidateVolumeCapabilitiesResponse{Message: ""}, nil
+}
+
+// ControllerGetCapabilities returns the capabilities of the Controller plugin
+func (d *DriverV2) ControllerGetCapabilities(ctx context.Context, req *csi.ControllerGetCapabilitiesRequest) (*csi.ControllerGetCapabilitiesResponse, error) {
+	return &csi.ControllerGetCapabilitiesResponse{
+		Capabilities: d.Cap,
+	}, nil
+}
+
+// GetCapacity returns the capacity of the total available storage pool
+func (d *DriverV2) GetCapacity(ctx context.Context, req *csi.GetCapacityRequest) (*csi.GetCapacityResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "")
+}
+
+// ListVolumes return all available volumes
+func (d *DriverV2) ListVolumes(ctx context.Context, req *csi.ListVolumesRequest) (*csi.ListVolumesResponse, error) {
+	start := 0
+	if req.StartingToken != "" {
+		var err error
+		start, err = strconv.Atoi(req.StartingToken)
+		if err != nil {
+			return nil, status.Errorf(codes.Aborted, "ListVolumes starting token(%s) parsing with error: %v", req.StartingToken, err)
+		}
+		if start < 0 {
+			return nil, status.Errorf(codes.Aborted, "ListVolumes starting token(%d) can not be negative", start)
+		}
+	}
+	if d.cloud.KubeClient != nil && d.cloud.KubeClient.CoreV1() != nil && d.cloud.KubeClient.CoreV1().PersistentVolumes() != nil {
+		klog.V(6).Infof("List Volumes in Cluster:")
+		return d.listVolumesInCluster(ctx, start, int(req.MaxEntries))
+	}
+	klog.V(6).Infof("List Volumes in Node Resource Group: %s", d.cloud.ResourceGroup)
+	return d.listVolumesInNodeResourceGroup(ctx, start, int(req.MaxEntries))
+}
+
+// listVolumesInCluster is a helper function for ListVolumes used for when there is an available kubeclient
+func (d *DriverV2) listVolumesInCluster(ctx context.Context, start, maxEntries int) (*csi.ListVolumesResponse, error) {
+	pvList, err := d.cloud.KubeClient.CoreV1().PersistentVolumes().List(context.TODO(), metav1.ListOptions{})
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "ListVolumes failed while fetching PersistentVolumes List with error: %v", err.Error())
+	}
+
+	// get all resource groups and put them into a sorted slice
+	rgMap := make(map[string]bool)
+	volSet := make(map[string]bool)
+	for _, pv := range pvList.Items {
+		if pv.Spec.CSI != nil && pv.Spec.CSI.Driver == DriverName {
+			diskURI := pv.Spec.CSI.VolumeHandle
+			if err := isValidDiskURI(diskURI); err != nil {
+				klog.Warningf("invalid disk uri (%s) with error(%v)", diskURI, err)
+				continue
+			}
+			rg, err := GetResourceGroupFromURI(diskURI)
+			if err != nil {
+				klog.Warningf("failed to get resource group from disk uri (%s) with error(%v)", diskURI, err)
+				continue
+			}
+			rg, diskURI = strings.ToLower(rg), strings.ToLower(diskURI)
+			volSet[diskURI] = true
+			if _, visited := rgMap[rg]; visited {
+				continue
+			}
+			rgMap[rg] = true
+		}
+	}
+
+	resourceGroups := make([]string, len(rgMap))
+	i := 0
+	for rg := range rgMap {
+		resourceGroups[i] = rg
+		i++
+	}
+	sort.Strings(resourceGroups)
+
+	// loop through each resourceGroup to get disk lists
+	entries := []*csi.ListVolumesResponse_Entry{}
+	numVisited := 0
+	isCompleteRun, startFound := true, false
+	for _, resourceGroup := range resourceGroups {
+		if !isCompleteRun || (maxEntries > 0 && len(entries) >= maxEntries) {
+			isCompleteRun = false
+			break
+		}
+		localStart := start - numVisited
+		if startFound {
+			localStart = 0
+		}
+		listStatus := d.listVolumesByResourceGroup(ctx, resourceGroup, entries, localStart, maxEntries-len(entries), volSet)
+		numVisited += listStatus.numVisited
+		if listStatus.err != nil {
+			if status.Code(listStatus.err) == codes.FailedPrecondition {
+				continue
+			}
+			return nil, listStatus.err
+		}
+		startFound = true
+		entries = listStatus.entries
+		isCompleteRun = isCompleteRun && listStatus.isCompleteRun
+	}
+	// if start was not found, start token was greater than total number of disks
+	if !startFound {
+		return nil, status.Errorf(codes.FailedPrecondition, "ListVolumes starting token(%d) is greater than total number of disks", start)
+	}
+
+	nextTokenString := ""
+	if !isCompleteRun {
+		nextTokenString = strconv.Itoa(start + numVisited)
+	}
+
+	listVolumesResp := &csi.ListVolumesResponse{
+		Entries:   entries,
+		NextToken: nextTokenString,
+	}
+
+	return listVolumesResp, nil
+}
+
+// listVolumesInNodeResourceGroup is a helper function for ListVolumes used for when there is no available kubeclient
+func (d *DriverV2) listVolumesInNodeResourceGroup(ctx context.Context, start, maxEntries int) (*csi.ListVolumesResponse, error) {
+	entries := []*csi.ListVolumesResponse_Entry{}
+	listStatus := d.listVolumesByResourceGroup(ctx, d.cloud.ResourceGroup, entries, start, maxEntries, nil)
+	if listStatus.err != nil {
+		return nil, listStatus.err
+	}
+
+	nextTokenString := ""
+	if !listStatus.isCompleteRun {
+		nextTokenString = strconv.Itoa(listStatus.numVisited)
+	}
+
+	listVolumesResp := &csi.ListVolumesResponse{
+		Entries:   listStatus.entries,
+		NextToken: nextTokenString,
+	}
+
+	return listVolumesResp, nil
+}
+
+// listVolumesByResourceGroup is a helper function that updates the ListVolumeResponse_Entry slice and returns number of total visited volumes, number of volumes that needs to be visited and an error if found
+func (d *DriverV2) listVolumesByResourceGroup(ctx context.Context, resourceGroup string, entries []*csi.ListVolumesResponse_Entry, start, maxEntries int, volSet map[string]bool) listVolumeStatus {
+	disks, derr := d.cloud.DisksClient.ListByResourceGroup(ctx, resourceGroup)
+	if derr != nil {
+		return listVolumeStatus{err: status.Errorf(codes.Internal, "ListVolumes on rg(%s) failed with error: %v", resourceGroup, derr.Error())}
+	}
+	// if volSet is initialized but is empty, return
+	if volSet != nil && len(volSet) == 0 {
+		return listVolumeStatus{
+			numVisited:    len(disks),
+			isCompleteRun: true,
+			entries:       entries,
+		}
+	}
+	if start > 0 && start >= len(disks) {
+		return listVolumeStatus{
+			numVisited: len(disks),
+			err:        status.Errorf(codes.FailedPrecondition, "ListVolumes starting token(%d) on rg(%s) is greater than total number of volumes", start, d.cloud.ResourceGroup),
+		}
+	}
+	if start < 0 {
+		start = 0
+	}
+	i := start
+	isCompleteRun := true
+	// Loop until
+	for ; i < len(disks); i++ {
+		if maxEntries > 0 && len(entries) >= maxEntries {
+			isCompleteRun = false
+			break
+		}
+
+		disk := disks[i]
+		// if given a set of volumes from KubeClient, only continue if the disk can be found in the set
+		if volSet != nil && !volSet[strings.ToLower(*disk.ID)] {
+			continue
+		}
+		// HyperVGeneration property is only setup for os disks. Only the non os disks should be included in the list
+		if disk.DiskProperties == nil || disk.DiskProperties.HyperVGeneration == "" {
+			nodeList := []string{}
+
+			if disk.ManagedBy != nil {
+				attachedNode, err := d.cloud.VMSet.GetNodeNameByProviderID(*disk.ManagedBy)
+				if err != nil {
+					return listVolumeStatus{err: err}
+				}
+				nodeList = append(nodeList, string(attachedNode))
+			}
+
+			entries = append(entries, &csi.ListVolumesResponse_Entry{
+				Volume: &csi.Volume{
+					VolumeId: *disk.ID,
+				},
+				Status: &csi.ListVolumesResponse_VolumeStatus{
+					PublishedNodeIds: nodeList,
+				},
+			})
+		}
+	}
+	return listVolumeStatus{
+		numVisited:    i - start,
+		isCompleteRun: isCompleteRun,
+		entries:       entries,
+	}
+}
+
+// ControllerExpandVolume controller expand volume
+func (d *DriverV2) ControllerExpandVolume(ctx context.Context, req *csi.ControllerExpandVolumeRequest) (*csi.ControllerExpandVolumeResponse, error) {
+	if len(req.GetVolumeId()) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
+	}
+	if err := d.ValidateControllerServiceRequest(csi.ControllerServiceCapability_RPC_EXPAND_VOLUME); err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid expand volume request: %v", req)
+	}
+
+	capacityBytes := req.GetCapacityRange().GetRequiredBytes()
+	if capacityBytes == 0 {
+		return nil, status.Error(codes.InvalidArgument, "volume capacity range missing in request")
+	}
+	requestSize := *resource.NewQuantity(capacityBytes, resource.BinarySI)
+
+	diskURI := req.GetVolumeId()
+	if err := isValidDiskURI(diskURI); err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "disk URI(%s) is not valid: %v", diskURI, err)
+	}
+
+	diskName, err := GetDiskName(diskURI)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "could not get disk name from diskURI(%s) with error(%v)", diskURI, err)
+	}
+	resourceGroup, err := GetResourceGroupFromURI(diskURI)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "could not get resource group from diskURI(%s) with error(%v)", diskURI, err)
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_expand_volume", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	result, rerr := d.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
+	if rerr != nil {
+		return nil, status.Errorf(codes.Internal, "could not get the disk(%s) under rg(%s) with error(%v)", diskName, resourceGroup, rerr.Error())
+	}
+	if result.DiskProperties.DiskSizeGB == nil {
+		return nil, status.Errorf(codes.Internal, "could not get size of the disk(%s)", diskName)
+	}
+	oldSize := *resource.NewQuantity(int64(*result.DiskProperties.DiskSizeGB), resource.BinarySI)
+
+	klog.V(2).Infof("begin to expand azure disk(%s) with new size(%d)", diskURI, requestSize.Value())
+	newSize, err := d.cloud.ResizeDisk(diskURI, oldSize, requestSize)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to resize disk(%s) with error(%v)", diskURI, err)
+	}
+
+	currentSize, ok := newSize.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform disk size with error(%v)", err)
+	}
+
+	isOperationSucceeded = true
+	klog.V(2).Infof("expand azure disk(%s) successfully, currentSize(%d)", diskURI, currentSize)
+
+	return &csi.ControllerExpandVolumeResponse{
+		CapacityBytes:         currentSize,
+		NodeExpansionRequired: true,
+	}, nil
+}
+
+// CreateSnapshot create a snapshot
+func (d *DriverV2) CreateSnapshot(ctx context.Context, req *csi.CreateSnapshotRequest) (*csi.CreateSnapshotResponse, error) {
+	sourceVolumeID := req.GetSourceVolumeId()
+	if len(sourceVolumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "CreateSnapshot Source Volume ID must be provided")
+	}
+	snapshotName := req.Name
+	if len(snapshotName) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "snapshot name must be provided")
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(sourceVolumeID); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, sourceVolumeID)
+	}
+	defer d.volumeLocks.Release(sourceVolumeID)
+
+	snapshotName = getValidDiskName(snapshotName)
+
+	var customTags string
+	// set incremental snapshot as true by default
+	incremental := true
+	var resourceGroup string
+	var err error
+
+	parameters := req.GetParameters()
+	for k, v := range parameters {
+		switch strings.ToLower(k) {
+		case tagsField:
+			customTags = v
+		case incrementalField:
+			if v == "false" {
+				incremental = false
+			}
+		case resourceGroupField:
+			resourceGroup = v
+		default:
+			return nil, fmt.Errorf("AzureDisk - invalid option %s in VolumeSnapshotClass", k)
+		}
+	}
+
+	if IsAzureStackCloud(d.cloud.Config.Cloud, d.supportAzureStack) {
+		klog.V(2).Info("Use full snapshot instead as Azure Stack does not incremental snapshot.")
+		incremental = false
+	}
+
+	if resourceGroup == "" {
+		resourceGroup, err = GetResourceGroupFromURI(sourceVolumeID)
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "could not get resource group from diskURI(%s) with error(%v)", sourceVolumeID, err)
+		}
+	}
+
+	customTagsMap, err := volumehelper.ConvertTagsToMap(customTags)
+	if err != nil {
+		return nil, err
+	}
+	tags := make(map[string]*string)
+	for k, v := range customTagsMap {
+		tags[k] = &v
+	}
+
+	snapshot := compute.Snapshot{
+		SnapshotProperties: &compute.SnapshotProperties{
+			CreationData: &compute.CreationData{
+				CreateOption: compute.Copy,
+				SourceURI:    &sourceVolumeID,
+			},
+			Incremental: &incremental,
+		},
+		Location: &d.cloud.Location,
+		Tags:     tags,
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_create_snapshot", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("begin to create snapshot(%s, incremental: %v) under rg(%s)", snapshotName, incremental, resourceGroup)
+	rerr := d.cloud.SnapshotsClient.CreateOrUpdate(ctx, resourceGroup, snapshotName, snapshot)
+	if rerr != nil {
+		if strings.Contains(rerr.Error().Error(), "existing disk") {
+			return nil, status.Error(codes.AlreadyExists, fmt.Sprintf("request snapshot(%s) under rg(%s) already exists, but the SourceVolumeId is different, error details: %v", snapshotName, resourceGroup, rerr.Error()))
+		}
+
+		return nil, status.Error(codes.Internal, fmt.Sprintf("create snapshot error: %v", rerr.Error()))
+	}
+	klog.V(2).Infof("create snapshot(%s) under rg(%s) successfully", snapshotName, resourceGroup)
+
+	csiSnapshot, err := d.getSnapshotByID(ctx, resourceGroup, snapshotName, sourceVolumeID)
+	if err != nil {
+		return nil, err
+	}
+
+	createResp := &csi.CreateSnapshotResponse{
+		Snapshot: csiSnapshot,
+	}
+	isOperationSucceeded = true
+	return createResp, nil
+}
+
+// DeleteSnapshot delete a snapshot
+func (d *DriverV2) DeleteSnapshot(ctx context.Context, req *csi.DeleteSnapshotRequest) (*csi.DeleteSnapshotResponse, error) {
+	snapshotID := req.SnapshotId
+	if len(snapshotID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Snapshot ID must be provided")
+	}
+
+	var err error
+	snapshotName := snapshotID
+	resourceGroup := d.cloud.ResourceGroup
+
+	if isARMResourceID(snapshotID) {
+		snapshotName, resourceGroup, err = d.getSnapshotInfo(snapshotID)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	mc := metrics.NewMetricContext(azureDiskCSIDriverName, "controller_delete_snapshot", d.cloud.ResourceGroup, d.cloud.SubscriptionID, d.Name)
+	isOperationSucceeded := false
+	defer func() {
+		mc.ObserveOperationWithResult(isOperationSucceeded)
+	}()
+
+	klog.V(2).Infof("begin to delete snapshot(%s) under rg(%s)", snapshotName, resourceGroup)
+	rerr := d.cloud.SnapshotsClient.Delete(ctx, resourceGroup, snapshotName)
+	if rerr != nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("delete snapshot error: %v", rerr.Error()))
+	}
+	klog.V(2).Infof("delete snapshot(%s) under rg(%s) successfully", snapshotName, resourceGroup)
+	isOperationSucceeded = true
+	return &csi.DeleteSnapshotResponse{}, nil
+}
+
+// ListSnapshots list all snapshots
+func (d *DriverV2) ListSnapshots(ctx context.Context, req *csi.ListSnapshotsRequest) (*csi.ListSnapshotsResponse, error) {
+	// SnapshotId is not empty, return snapshot that match the snapshot id.
+	if len(req.GetSnapshotId()) != 0 {
+		snapshot, err := d.getSnapshotByID(ctx, d.cloud.ResourceGroup, req.GetSnapshotId(), req.SourceVolumeId)
+		if err != nil {
+			if strings.Contains(err.Error(), resourceNotFound) {
+				return &csi.ListSnapshotsResponse{}, nil
+			}
+			return nil, err
+		}
+		entries := []*csi.ListSnapshotsResponse_Entry{
+			{
+				Snapshot: snapshot,
+			},
+		}
+		listSnapshotResp := &csi.ListSnapshotsResponse{
+			Entries: entries,
+		}
+		return listSnapshotResp, nil
+	}
+
+	// no SnapshotId is set, return all snapshots that satisfy the request.
+	snapshots, err := d.cloud.SnapshotsClient.ListByResourceGroup(ctx, d.cloud.ResourceGroup)
+	if err != nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("Unknown list snapshot error: %v", err.Error()))
+	}
+
+	return getEntriesAndNextToken(req, snapshots)
+}
+
+func (d *DriverV2) getSnapshotByID(ctx context.Context, resourceGroup, snapshotID, sourceVolumeID string) (*csi.Snapshot, error) {
+	var err error
+	snapshotName := snapshotID
+	if isARMResourceID(snapshotID) {
+		snapshotName, resourceGroup, err = d.getSnapshotInfo(snapshotID)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	snapshot, rerr := d.cloud.SnapshotsClient.Get(ctx, resourceGroup, snapshotName)
+	if rerr != nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("get snapshot %s from rg(%s) error: %v", snapshotName, resourceGroup, rerr.Error()))
+	}
+
+	return generateCSISnapshot(sourceVolumeID, &snapshot)
+}
+
+// GetSourceDiskSize recursively searches for the sourceDisk and returns: sourceDisk disk size, error
+func (d *DriverV2) GetSourceDiskSize(ctx context.Context, resourceGroup, diskName string, curDepth, maxDepth int) (*int32, error) {
+	if curDepth > maxDepth {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("current depth (%d) surpassed the max depth (%d) while searching for the source disk size", curDepth, maxDepth))
+	}
+	result, rerr := d.cloud.DisksClient.Get(ctx, resourceGroup, diskName)
+	if rerr != nil {
+		return nil, rerr.Error()
+	}
+	if result.DiskProperties == nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("DiskProperty not found for disk (%s) in resource group (%s)", diskName, resourceGroup))
+	}
+
+	if result.DiskProperties.CreationData != nil && (*result.DiskProperties.CreationData).CreateOption == "Copy" {
+		klog.V(2).Infof("Clone source disk has a parent source")
+		sourceResourceID := *result.DiskProperties.CreationData.SourceResourceID
+		parentResourceGroup, _ := GetResourceGroupFromURI(sourceResourceID)
+		parentDiskName := path.Base(sourceResourceID)
+		return d.GetSourceDiskSize(ctx, parentResourceGroup, parentDiskName, curDepth+1, maxDepth)
+	}
+
+	if (*result.DiskProperties).DiskSizeGB == nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("DiskSizeGB for disk (%s) in resourcegroup (%s) is nil", diskName, resourceGroup))
+	}
+	return (*result.DiskProperties).DiskSizeGB, nil
+}
+
+// The format of snapshot id is /subscriptions/xxx/resourceGroups/xxx/providers/Microsoft.Compute/snapshots/snapshot-xxx-xxx.
+func (d *DriverV2) getSnapshotInfo(snapshotID string) (snapshotName, resourceGroup string, err error) {
+	if snapshotName, err = getSnapshotName(snapshotID); err != nil {
+		return "", "", err
+	}
+	if resourceGroup, err = GetResourceGroupFromURI(snapshotID); err != nil {
+		return "", "", err
+	}
+	return snapshotName, resourceGroup, err
+}
diff --git a/pkg/azuredisk/fake_azuredisk.go b/pkg/azuredisk/fake_azuredisk.go
index c664500a1..9b1bccf2d 100644
--- a/pkg/azuredisk/fake_azuredisk.go
+++ b/pkg/azuredisk/fake_azuredisk.go
@@ -17,14 +17,17 @@ limitations under the License.
 package azuredisk
 
 import (
+	"context"
 	"testing"
 
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/golang/mock/gomock"
+	"k8s.io/mount-utils"
 
 	csicommon "sigs.k8s.io/azuredisk-csi-driver/pkg/csi-common"
 	"sigs.k8s.io/azuredisk-csi-driver/pkg/mounter"
 	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+	"sigs.k8s.io/cloud-provider-azure/pkg/provider"
 	azure "sigs.k8s.io/cloud-provider-azure/pkg/provider"
 )
 
@@ -52,7 +55,31 @@ var (
 	}
 )
 
-func NewFakeDriver(t *testing.T) (*Driver, error) {
+// FakeDriver defines an interface unit tests use to test either the v1 or v2 implementation of the Azure Disk CSI Driver.
+type FakeDriver interface {
+	CSIDriver
+
+	GetSourceDiskSize(ctx context.Context, resourceGroup, diskName string, curDepth, maxDepth int) (*int32, error)
+
+	setControllerCapabilities([]*csi.ControllerServiceCapability)
+	setNodeCapabilities([]*csi.NodeServiceCapability)
+	setName(string)
+	setNodeID(string)
+	setVersion(version string)
+	getCloud() *provider.Cloud
+	setCloud(*provider.Cloud)
+	getMounter() *mount.SafeFormatAndMount
+	setMounter(*mount.SafeFormatAndMount)
+
+	checkDiskCapacity(context.Context, string, string, int) (bool, error)
+	getSnapshotInfo(string) (string, string, error)
+	getSnapshotByID(context.Context, string, string, string) (*csi.Snapshot, error)
+	ensureMountPoint(string) (bool, error)
+	ensureBlockTargetFile(string) error
+	getDevicePathWithLUN(lunStr string) (string, error)
+}
+
+func newFakeDriverV1(t *testing.T) (*Driver, error) {
 	driver := Driver{}
 	driver.Name = fakeDriverName
 	driver.Version = fakeDriverVersion
diff --git a/pkg/azuredisk/fake_azuredisk_v1.go b/pkg/azuredisk/fake_azuredisk_v1.go
new file mode 100644
index 000000000..5ed0124a1
--- /dev/null
+++ b/pkg/azuredisk/fake_azuredisk_v1.go
@@ -0,0 +1,26 @@
+// +build !azurediskv2
+
+/*
+Copyright 2020 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import "testing"
+
+// NewFakeDriver returns a driver implementation suitable for use in unit tests.
+func NewFakeDriver(t *testing.T) (FakeDriver, error) {
+	return newFakeDriverV1(t)
+}
diff --git a/pkg/azuredisk/fake_azuredisk_v2.go b/pkg/azuredisk/fake_azuredisk_v2.go
new file mode 100644
index 000000000..9c69fc93b
--- /dev/null
+++ b/pkg/azuredisk/fake_azuredisk_v2.go
@@ -0,0 +1,85 @@
+// +build azurediskv2
+
+/*
+Copyright 2020 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import (
+	"testing"
+
+	"github.com/container-storage-interface/spec/lib/go/csi"
+	"github.com/golang/mock/gomock"
+	"k8s.io/klog/v2"
+	csicommon "sigs.k8s.io/azuredisk-csi-driver/pkg/csi-common"
+	"sigs.k8s.io/azuredisk-csi-driver/pkg/mounter"
+	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+	azure "sigs.k8s.io/cloud-provider-azure/pkg/provider"
+)
+
+// NewFakeDriver returns a driver implementation suitable for use in unit tests.
+func NewFakeDriver(t *testing.T) (FakeDriver, error) {
+	var d FakeDriver
+	var err error
+
+	if !*useDriverV2 {
+		d, err = newFakeDriverV1(t)
+	} else {
+		d, err = newFakeDriverV2(t)
+	}
+
+	return d, err
+}
+
+func newFakeDriverV2(t *testing.T) (*DriverV2, error) {
+	klog.Warning("Using DriverV2")
+	driver := DriverV2{}
+	driver.Name = fakeDriverName
+	driver.Version = fakeDriverVersion
+	driver.NodeID = fakeNodeID
+	driver.CSIDriver = *csicommon.NewFakeCSIDriver()
+	driver.volumeLocks = volumehelper.NewVolumeLocks()
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	driver.cloud = azure.GetTestCloud(ctrl)
+	mounter, err := mounter.NewSafeMounter()
+	if err != nil {
+		return nil, err
+	}
+
+	driver.mounter = mounter
+
+	driver.AddControllerServiceCapabilities(
+		[]csi.ControllerServiceCapability_RPC_Type{
+			csi.ControllerServiceCapability_RPC_CREATE_DELETE_VOLUME,
+			csi.ControllerServiceCapability_RPC_PUBLISH_UNPUBLISH_VOLUME,
+			csi.ControllerServiceCapability_RPC_CREATE_DELETE_SNAPSHOT,
+			csi.ControllerServiceCapability_RPC_LIST_SNAPSHOTS,
+			csi.ControllerServiceCapability_RPC_CLONE_VOLUME,
+			csi.ControllerServiceCapability_RPC_EXPAND_VOLUME,
+			csi.ControllerServiceCapability_RPC_LIST_VOLUMES,
+			csi.ControllerServiceCapability_RPC_LIST_VOLUMES_PUBLISHED_NODES,
+		})
+	driver.AddVolumeCapabilityAccessModes([]csi.VolumeCapability_AccessMode_Mode{csi.VolumeCapability_AccessMode_SINGLE_NODE_WRITER})
+	driver.AddNodeServiceCapabilities([]csi.NodeServiceCapability_RPC_Type{
+		csi.NodeServiceCapability_RPC_STAGE_UNSTAGE_VOLUME,
+		csi.NodeServiceCapability_RPC_EXPAND_VOLUME,
+	})
+
+	return &driver, nil
+}
diff --git a/pkg/azuredisk/identityserver_v2.go b/pkg/azuredisk/identityserver_v2.go
new file mode 100644
index 000000000..4c50b7858
--- /dev/null
+++ b/pkg/azuredisk/identityserver_v2.go
@@ -0,0 +1,76 @@
+// +build azurediskv2
+
+/*
+Copyright 2017 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/container-storage-interface/spec/lib/go/csi"
+
+	wrappers "github.com/golang/protobuf/ptypes/wrappers"
+)
+
+// GetPluginInfo return the version and name of the plugin
+func (f *DriverV2) GetPluginInfo(ctx context.Context, req *csi.GetPluginInfoRequest) (*csi.GetPluginInfoResponse, error) {
+	if f.Name == "" {
+		return nil, status.Error(codes.Unavailable, "DriverV2 name not configured")
+	}
+
+	if f.Version == "" {
+		return nil, status.Error(codes.Unavailable, "DriverV2 is missing version")
+	}
+
+	return &csi.GetPluginInfoResponse{
+		Name:          f.Name,
+		VendorVersion: f.Version,
+	}, nil
+}
+
+// Probe check whether the plugin is running or not.
+// This method does not need to return anything.
+// Currently the spec does not dictate what you should return either.
+// Hence, return an empty response
+func (f *DriverV2) Probe(ctx context.Context, req *csi.ProbeRequest) (*csi.ProbeResponse, error) {
+	return &csi.ProbeResponse{Ready: &wrappers.BoolValue{Value: true}}, nil
+}
+
+// GetPluginCapabilities returns the capabilities of the plugin
+func (f *DriverV2) GetPluginCapabilities(ctx context.Context, req *csi.GetPluginCapabilitiesRequest) (*csi.GetPluginCapabilitiesResponse, error) {
+	return &csi.GetPluginCapabilitiesResponse{
+		Capabilities: []*csi.PluginCapability{
+			{
+				Type: &csi.PluginCapability_Service_{
+					Service: &csi.PluginCapability_Service{
+						Type: csi.PluginCapability_Service_CONTROLLER_SERVICE,
+					},
+				},
+			},
+			{
+				Type: &csi.PluginCapability_Service_{
+					Service: &csi.PluginCapability_Service{
+						Type: csi.PluginCapability_Service_VOLUME_ACCESSIBILITY_CONSTRAINTS,
+					},
+				},
+			},
+		},
+	}, nil
+}
diff --git a/pkg/azuredisk/nodeserver_v2.go b/pkg/azuredisk/nodeserver_v2.go
new file mode 100644
index 000000000..daf40f0e1
--- /dev/null
+++ b/pkg/azuredisk/nodeserver_v2.go
@@ -0,0 +1,551 @@
+// +build azurediskv2
+
+/*
+Copyright 2017 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package azuredisk
+
+import (
+	"context"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"runtime"
+	"time"
+
+	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
+
+	"github.com/container-storage-interface/spec/lib/go/csi"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/wait"
+	"k8s.io/klog/v2"
+	"k8s.io/kubernetes/pkg/util/resizefs"
+	"k8s.io/kubernetes/pkg/volume"
+	"k8s.io/kubernetes/pkg/volume/util/hostutil"
+)
+
+// NodeStageVolume mount disk device to a staging path
+func (d *DriverV2) NodeStageVolume(ctx context.Context, req *csi.NodeStageVolumeRequest) (*csi.NodeStageVolumeResponse, error) {
+	diskURI := req.GetVolumeId()
+	if len(diskURI) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID not provided")
+	}
+
+	target := req.GetStagingTargetPath()
+	if len(target) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Staging target not provided")
+	}
+
+	volumeCapability := req.GetVolumeCapability()
+	if volumeCapability == nil {
+		return nil, status.Error(codes.InvalidArgument, "Volume capability not provided")
+	}
+
+	if !isValidVolumeCapabilities([]*csi.VolumeCapability{volumeCapability}) {
+		return nil, status.Error(codes.InvalidArgument, "Volume capability not supported")
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(diskURI); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, diskURI)
+	}
+	defer d.volumeLocks.Release(diskURI)
+
+	// If the access type is block, do nothing for stage
+	switch req.GetVolumeCapability().GetAccessType().(type) {
+	case *csi.VolumeCapability_Block:
+		return &csi.NodeStageVolumeResponse{}, nil
+	}
+
+	lun, ok := req.PublishContext[LUN]
+	if !ok {
+		return nil, status.Error(codes.InvalidArgument, "lun not provided")
+	}
+
+	mnt, err := d.ensureMountPoint(target)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not mount target %q: %v", target, err)
+	}
+	if mnt {
+		klog.V(2).Infof("NodeStageVolume: already mounted on target %s", target)
+		return &csi.NodeStageVolumeResponse{}, nil
+	}
+
+	// Get fsType and mountOptions that the volume will be formatted and mounted with
+	fstype := getDefaultFsType()
+	options := []string{}
+	if mnt := volumeCapability.GetMount(); mnt != nil {
+		if mnt.FsType != "" {
+			fstype = mnt.FsType
+		}
+		options = append(options, mnt.MountFlags...)
+	}
+
+	volContextFSType := getFStype(req.GetVolumeContext())
+	if volContextFSType != "" {
+		// respect "fstype" setting in storage class parameters
+		fstype = volContextFSType
+	}
+
+	source, err := d.getDevicePathWithLUN(lun)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "Failed to find disk on lun %s. %v", lun, err)
+	}
+
+	// If partition is specified, should mount it only instead of the entire disk.
+	if partition, ok := req.GetVolumeContext()[volumeAttributePartition]; ok {
+		source = source + "-part" + partition
+	}
+
+	// FormatAndMount will format only if needed
+	klog.V(2).Infof("NodeStageVolume: formatting %s and mounting at %s with mount options(%s)", source, target, options)
+	err = d.formatAndMount(source, target, fstype, options)
+	if err != nil {
+		msg := fmt.Sprintf("could not format %q(lun: %q), and mount it at %q", source, lun, target)
+		return nil, status.Error(codes.Internal, msg)
+	}
+	klog.V(2).Infof("NodeStageVolume: format %s and mounting at %s successfully.", source, target)
+
+	// if resize is required, resize filesystem
+	if required, ok := req.GetVolumeContext()[resizeRequired]; ok && required == "true" {
+		klog.V(2).Infof("NodeStageVolume: fs resize initiating on target(%s) volumeid(%s)", target, diskURI)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "NodeStageVolume: Could not get volume path for %s: %v", target, err)
+		}
+
+		resizer := resizefs.NewResizeFs(d.mounter)
+		if _, err := resizer.Resize(source, target); err != nil {
+			return nil, status.Errorf(codes.Internal, "NodeStageVolume: Could not resize volume %q (%q):  %v", diskURI, source, err)
+		}
+
+		klog.V(2).Infof("NodeStageVolume: fs resize successful on target(%s) volumeid(%s).", target, diskURI)
+	}
+
+	return &csi.NodeStageVolumeResponse{}, nil
+}
+
+// NodeUnstageVolume unmount disk device from a staging path
+func (d *DriverV2) NodeUnstageVolume(ctx context.Context, req *csi.NodeUnstageVolumeRequest) (*csi.NodeUnstageVolumeResponse, error) {
+	volumeID := req.GetVolumeId()
+	if len(volumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID not provided")
+	}
+
+	stagingTargetPath := req.GetStagingTargetPath()
+	if len(stagingTargetPath) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Staging target not provided")
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(volumeID); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, volumeID)
+	}
+	defer d.volumeLocks.Release(volumeID)
+
+	klog.V(2).Infof("NodeUnstageVolume: unmounting %s", stagingTargetPath)
+	err := CleanupMountPoint(stagingTargetPath, d.mounter, false)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to unmount staging target %q: %v", stagingTargetPath, err)
+	}
+	klog.V(2).Infof("NodeUnstageVolume: unmount %s successfully", stagingTargetPath)
+
+	return &csi.NodeUnstageVolumeResponse{}, nil
+}
+
+// NodePublishVolume mount the volume from staging to target path
+func (d *DriverV2) NodePublishVolume(ctx context.Context, req *csi.NodePublishVolumeRequest) (*csi.NodePublishVolumeResponse, error) {
+	volumeID := req.GetVolumeId()
+	if req.GetVolumeCapability() == nil {
+		return nil, status.Error(codes.InvalidArgument, "Volume capability missing in request")
+	}
+	if len(volumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
+	}
+
+	source := req.GetStagingTargetPath()
+	if len(source) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Staging target not provided")
+	}
+
+	target := req.GetTargetPath()
+	if len(target) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Target path not provided")
+	}
+
+	err := preparePublishPath(target, d.mounter)
+	if err != nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("Target path could not be prepared: %v", err))
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(volumeID); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, volumeID)
+	}
+	defer d.volumeLocks.Release(volumeID)
+
+	mountOptions := []string{"bind"}
+	if req.GetReadonly() {
+		mountOptions = append(mountOptions, "ro")
+	}
+
+	switch req.GetVolumeCapability().GetAccessType().(type) {
+	case *csi.VolumeCapability_Block:
+		lun, ok := req.PublishContext[LUN]
+		if !ok {
+			return nil, status.Error(codes.InvalidArgument, "lun not provided")
+		}
+		var err error
+		source, err = d.getDevicePathWithLUN(lun)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "Failed to find device path with lun %s. %v", lun, err)
+		}
+		klog.V(2).Infof("NodePublishVolume [block]: found device path %s with lun %s", source, lun)
+		err = d.ensureBlockTargetFile(target)
+		if err != nil {
+			return nil, err
+		}
+	case *csi.VolumeCapability_Mount:
+		mnt, err := d.ensureMountPoint(target)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "Could not mount target %q: %v", target, err)
+		}
+		if mnt {
+			klog.V(2).Infof("NodePublishVolume: already mounted on target %s", target)
+			return &csi.NodePublishVolumeResponse{}, nil
+		}
+	}
+
+	klog.V(2).Infof("NodePublishVolume: mounting %s at %s", source, target)
+	if err := d.mounter.Mount(source, target, "", mountOptions); err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not mount %q at %q: %v", source, target, err)
+	}
+
+	klog.V(2).Infof("NodePublishVolume: mount %s at %s successfully", source, target)
+
+	return &csi.NodePublishVolumeResponse{}, nil
+}
+
+// NodeUnpublishVolume unmount the volume from the target path
+func (d *DriverV2) NodeUnpublishVolume(ctx context.Context, req *csi.NodeUnpublishVolumeRequest) (*csi.NodeUnpublishVolumeResponse, error) {
+	targetPath := req.GetTargetPath()
+	volumeID := req.GetVolumeId()
+
+	if len(volumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID missing in request")
+	}
+	if len(targetPath) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Target path missing in request")
+	}
+
+	klog.V(2).Infof("NodeUnpublishVolume: unmounting volume %s on %s", volumeID, targetPath)
+	err := CleanupMountPoint(targetPath, d.mounter, false)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "failed to unmount target %q: %v", targetPath, err)
+	}
+
+	if acquired := d.volumeLocks.TryAcquire(volumeID); !acquired {
+		return nil, status.Errorf(codes.Aborted, volumeOperationAlreadyExistsFmt, volumeID)
+	}
+	defer d.volumeLocks.Release(volumeID)
+
+	klog.V(2).Infof("NodeUnpublishVolume: unmount volume %s on %s successfully", volumeID, targetPath)
+
+	return &csi.NodeUnpublishVolumeResponse{}, nil
+}
+
+// NodeGetCapabilities return the capabilities of the Node plugin
+func (d *DriverV2) NodeGetCapabilities(ctx context.Context, req *csi.NodeGetCapabilitiesRequest) (*csi.NodeGetCapabilitiesResponse, error) {
+	return &csi.NodeGetCapabilitiesResponse{
+		Capabilities: d.NSCap,
+	}, nil
+}
+
+// NodeGetInfo return info of the node on which this plugin is running
+func (d *DriverV2) NodeGetInfo(ctx context.Context, req *csi.NodeGetInfoRequest) (*csi.NodeGetInfoResponse, error) {
+	instances, ok := d.cloud.Instances()
+	if !ok {
+		return nil, status.Error(codes.Internal, "Failed to get instances from cloud provider")
+	}
+
+	instanceType, err := instances.InstanceType(context.TODO(), types.NodeName(d.NodeID))
+	if err != nil {
+		klog.Warningf("Failed to get instance type from Azure cloud provider, nodeName: %v, error: %v", d.NodeID, err)
+		instanceType = ""
+	}
+
+	topology := &csi.Topology{
+		Segments: map[string]string{topologyKey: ""},
+	}
+	zone, err := d.cloud.GetZone(ctx)
+	if err != nil {
+		klog.Warningf("Failed to get zone from Azure cloud provider, nodeName: %v, error: %v", d.NodeID, err)
+	} else {
+		if isAvailabilityZone(zone.FailureDomain, d.cloud.Location) {
+			topology.Segments[topologyKey] = zone.FailureDomain
+			klog.V(2).Infof("NodeGetInfo, nodeName: %v, zone: %v", d.NodeID, zone.FailureDomain)
+		}
+	}
+
+	return &csi.NodeGetInfoResponse{
+		NodeId:             d.NodeID,
+		MaxVolumesPerNode:  getMaxDataDiskCount(instanceType),
+		AccessibleTopology: topology,
+	}, nil
+}
+
+func (d *DriverV2) NodeGetVolumeStats(ctx context.Context, req *csi.NodeGetVolumeStatsRequest) (*csi.NodeGetVolumeStatsResponse, error) {
+	if len(req.VolumeId) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "NodeGetVolumeStats volume ID was empty")
+	}
+	if len(req.VolumePath) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "NodeGetVolumeStats volume path was empty")
+	}
+
+	_, err := os.Stat(req.VolumePath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil, status.Errorf(codes.NotFound, "path %s does not exist", req.VolumePath)
+		}
+		return nil, status.Errorf(codes.Internal, "failed to stat file %s: %v", req.VolumePath, err)
+	}
+
+	isBlock, err := hostutil.NewHostUtil().PathIsDevice(req.VolumePath)
+	if err != nil {
+		return nil, status.Errorf(codes.NotFound, "failed to determine whether %s is block device: %v", req.VolumePath, err)
+	}
+	if isBlock {
+		bcap, err := getBlockSizeBytes(req.VolumePath, d.mounter)
+		if err != nil {
+			return nil, status.Errorf(codes.Internal, "failed to get block capacity on path %s: %v", req.VolumePath, err)
+		}
+		return &csi.NodeGetVolumeStatsResponse{
+			Usage: []*csi.VolumeUsage{
+				{
+					Unit:  csi.VolumeUsage_BYTES,
+					Total: bcap,
+				},
+			},
+		}, nil
+	}
+
+	volumeMetrics, err := volume.NewMetricsStatFS(req.VolumePath).GetMetrics()
+	if err != nil {
+		return nil, err
+	}
+
+	available, ok := volumeMetrics.Available.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform volume available size(%v)", volumeMetrics.Available)
+	}
+	capacity, ok := volumeMetrics.Capacity.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform volume capacity size(%v)", volumeMetrics.Capacity)
+	}
+	used, ok := volumeMetrics.Used.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform volume used size(%v)", volumeMetrics.Used)
+	}
+
+	inodesFree, ok := volumeMetrics.InodesFree.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform disk inodes free(%v)", volumeMetrics.InodesFree)
+	}
+	inodes, ok := volumeMetrics.Inodes.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform disk inodes(%v)", volumeMetrics.Inodes)
+	}
+	inodesUsed, ok := volumeMetrics.InodesUsed.AsInt64()
+	if !ok {
+		return nil, status.Errorf(codes.Internal, "failed to transform disk inodes used(%v)", volumeMetrics.InodesUsed)
+	}
+
+	return &csi.NodeGetVolumeStatsResponse{
+		Usage: []*csi.VolumeUsage{
+			{
+				Unit:      csi.VolumeUsage_BYTES,
+				Available: available,
+				Total:     capacity,
+				Used:      used,
+			},
+			{
+				Unit:      csi.VolumeUsage_INODES,
+				Available: inodesFree,
+				Total:     inodes,
+				Used:      inodesUsed,
+			},
+		},
+	}, nil
+}
+
+// NodeExpandVolume node expand volume
+func (d *DriverV2) NodeExpandVolume(ctx context.Context, req *csi.NodeExpandVolumeRequest) (*csi.NodeExpandVolumeResponse, error) {
+	volumeID := req.GetVolumeId()
+	if len(volumeID) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "Volume ID not provided")
+	}
+	capacityBytes := req.GetCapacityRange().GetRequiredBytes()
+	volSizeBytes := int64(capacityBytes)
+	requestGiB := volumehelper.RoundUpGiB(volSizeBytes)
+
+	volumePath := req.GetVolumePath()
+	if len(volumePath) == 0 {
+		return nil, status.Error(codes.InvalidArgument, "volume path must be provided")
+	}
+
+	isBlock, err := hostutil.NewHostUtil().PathIsDevice(volumePath)
+	if err != nil {
+		return nil, status.Errorf(codes.NotFound, "failed to determine device path for volumePath [%v]: %v", volumePath, err)
+	}
+	if isBlock {
+		// Noop for Block NodeExpandVolume
+		klog.V(4).Infof("NodeExpandVolume succeeded on %v to %s, path check is block so this is a no-op", volumeID, volumePath)
+		return &csi.NodeExpandVolumeResponse{}, nil
+	}
+
+	volumeCapability := req.GetVolumeCapability()
+	if volumeCapability != nil {
+		// VolumeCapability is optional, if specified, validate it
+		caps := []*csi.VolumeCapability{volumeCapability}
+		if !isValidVolumeCapabilities(caps) {
+			return nil, status.Error(codes.InvalidArgument, "VolumeCapability is invalid.")
+		}
+
+		if blk := volumeCapability.GetBlock(); blk != nil {
+			// Noop for Block NodeExpandVolume
+			// This should not be executed but if somehow it is set to Block we should be cautious
+			klog.Warningf("NodeExpandVolume succeeded on %v to %s, capability is block but block check failed to identify it", volumeID, volumePath)
+			return &csi.NodeExpandVolumeResponse{}, nil
+		}
+	}
+
+	devicePath, err := getDevicePathWithMountPath(volumePath, d.mounter)
+	if err != nil {
+		return nil, status.Errorf(codes.NotFound, err.Error())
+	}
+
+	if err := resizeVolume(devicePath, volumePath, d.mounter); err != nil {
+		return nil, status.Errorf(codes.Internal, "Could not resize volume %q (%q):  %v", volumeID, devicePath, err)
+	}
+
+	gotBlockSizeBytes, err := getBlockSizeBytes(devicePath, d.mounter)
+	if err != nil {
+		return nil, status.Error(codes.Internal, fmt.Sprintf("Could not get size of block volume at path %s: %v", devicePath, err))
+	}
+	gotBlockGiB := volumehelper.RoundUpGiB(gotBlockSizeBytes)
+	if gotBlockGiB < requestGiB {
+		// Because size was rounded up, getting more size than requested will be a success.
+		return nil, status.Errorf(codes.Internal, "resize requested for %v, but after resizing volume size was %v", requestGiB, gotBlockGiB)
+	}
+	klog.V(2).Infof("NodeExpandVolume succeeded on resizing volume %v to %v", volumeID, gotBlockSizeBytes)
+
+	return &csi.NodeExpandVolumeResponse{
+		CapacityBytes: gotBlockSizeBytes,
+	}, nil
+}
+
+// ensureMountPoint: create mount point if not exists
+// return <true, nil> if it's already a mounted point otherwise return <false, nil>
+func (d *DriverV2) ensureMountPoint(target string) (bool, error) {
+	notMnt, err := d.mounter.IsLikelyNotMountPoint(target)
+	if err != nil && !os.IsNotExist(err) {
+		if IsCorruptedDir(target) {
+			notMnt = false
+			klog.Warningf("detected corrupted mount for targetPath [%s]", target)
+		} else {
+			return !notMnt, err
+		}
+	}
+
+	if !notMnt {
+		// testing original mount point, make sure the mount link is valid
+		_, err := ioutil.ReadDir(target)
+		if err == nil {
+			klog.V(2).Infof("already mounted to target %s", target)
+			return !notMnt, nil
+		}
+		// mount link is invalid, now unmount and remount later
+		klog.Warningf("ReadDir %s failed with %v, unmount this directory", target, err)
+		if err := d.mounter.Unmount(target); err != nil {
+			klog.Errorf("Unmount directory %s failed with %v", target, err)
+			return !notMnt, err
+		}
+		notMnt = true
+		return !notMnt, err
+	}
+
+	if runtime.GOOS != "windows" {
+		// in windows, we will use mklink to mount, will MkdirAll in Mount func
+		if err := volumehelper.MakeDir(target); err != nil {
+			klog.Errorf("mkdir failed on target: %s (%v)", target, err)
+			return !notMnt, err
+		}
+	}
+
+	return !notMnt, nil
+}
+
+func (d *DriverV2) formatAndMount(source, target, fstype string, options []string) error {
+	return formatAndMount(source, target, fstype, options, d.mounter)
+}
+
+func (d *DriverV2) getDevicePathWithLUN(lunStr string) (string, error) {
+	lun, err := getDiskLUN(lunStr)
+	if err != nil {
+		return "", err
+	}
+
+	io := &osIOHandler{}
+	scsiHostRescan(io, d.mounter)
+
+	newDevicePath := ""
+	err = wait.PollImmediate(1*time.Second, 2*time.Minute, func() (bool, error) {
+		var err error
+		if newDevicePath, err = findDiskByLun(int(lun), io, d.mounter); err != nil {
+			return false, fmt.Errorf("azureDisk - findDiskByLun(%v) failed with error(%s)", lun, err)
+		}
+
+		// did we find it?
+		if newDevicePath != "" {
+			return true, nil
+		}
+		// wait until timeout
+		return false, nil
+	})
+	if err == nil && newDevicePath == "" {
+		err = fmt.Errorf("azureDisk - findDiskByLun(%v) failed within timeout", lun)
+	}
+	return newDevicePath, err
+}
+
+func (d *DriverV2) ensureBlockTargetFile(target string) error {
+	// Since the block device target path is file, its parent directory should be ensured to be valid.
+	parentDir := filepath.Dir(target)
+	if _, err := d.ensureMountPoint(parentDir); err != nil {
+		return status.Errorf(codes.Internal, "Could not mount target %q: %v", parentDir, err)
+	}
+	// Create the mount point as a file since bind mount device node requires it to be a file
+	klog.V(2).Infof("ensureBlockTargetFile [block]: making target file %s", target)
+	err := volumehelper.MakeFile(target)
+	if err != nil {
+		if removeErr := os.Remove(target); removeErr != nil {
+			return status.Errorf(codes.Internal, "Could not remove mount target %q: %v", target, removeErr)
+		}
+		return status.Errorf(codes.Internal, "Could not create file %q: %v", target, err)
+	}
+
+	return nil
+}


Test patch:
diff --git a/pkg/azuredisk/azuredisk_test.go b/pkg/azuredisk/azuredisk_test.go
index 3fc550b97..75dc3a260 100644
--- a/pkg/azuredisk/azuredisk_test.go
+++ b/pkg/azuredisk/azuredisk_test.go
@@ -31,7 +31,6 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/assert"
 
-	"sigs.k8s.io/azuredisk-csi-driver/pkg/util"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azureclients/diskclient/mockdiskclient"
 	azure "sigs.k8s.io/cloud-provider-azure/pkg/provider"
 )
@@ -593,18 +592,6 @@ func TestIsCorruptedDir(t *testing.T) {
 	}
 }
 
-func TestNewDriver(t *testing.T) {
-	nodeid := "test"
-	driver := Driver{}
-	driver.Name = DriverName
-	driver.Version = driverVersion
-	driver.NodeID = nodeid
-	driver.volumeLocks = util.NewVolumeLocks()
-	driver.supportAzureStack = true
-	newdriver := NewDriver(nodeid, true)
-	assert.Equal(t, driver, *newdriver)
-}
-
 func TestCheckDiskExists(t *testing.T) {
 	diskName := "-"
 	assert.Equal(t, false, checkDiskName(diskName))
@@ -622,7 +609,7 @@ func TestCheckDiskCapacity(t *testing.T) {
 			DiskSizeGB: &size,
 		},
 	}
-	d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+	d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 	flag, err := d.checkDiskCapacity(context.TODO(), resourceGroup, diskName, 10)
 	assert.Equal(t, flag, true)
 	assert.Nil(t, err)
@@ -696,8 +683,8 @@ func TestRun(t *testing.T) {
 				os.Setenv(DefaultAzureCredentialFileEnv, fakeCredFile)
 
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
-				d.NodeID = ""
+				d.setCloud(&azure.Cloud{})
+				d.setNodeID("")
 				d.Run("tcp://127.0.0.1:0", "", true)
 			},
 		},
diff --git a/pkg/azuredisk/controllerserver_test.go b/pkg/azuredisk/controllerserver_test.go
index 371073f2b..1b25817fd 100644
--- a/pkg/azuredisk/controllerserver_test.go
+++ b/pkg/azuredisk/controllerserver_test.go
@@ -205,7 +205,7 @@ func TestCreateVolume(t *testing.T) {
 			name: " invalid ",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.Cap = []*csi.ControllerServiceCapability{}
+				d.setControllerCapabilities([]*csi.ControllerServiceCapability{})
 
 				req := &csi.CreateVolumeRequest{}
 				_, err := d.CreateVolume(context.Background(), req)
@@ -375,7 +375,7 @@ func TestCreateVolume(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &compute.DiskProperties{},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				_, err := d.CreateVolume(context.Background(), req)
 				expectedErr := fmt.Errorf("Tags 'unit-test' are invalid, the format should like: 'key1=value1,key2=value2'")
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -407,11 +407,11 @@ func TestCreateVolume(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &compute.DiskProperties{},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("test"),
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(rerr).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(rerr).AnyTimes()
 				_, err := d.CreateVolume(context.Background(), req)
 				expectedErr := fmt.Errorf("Retriable: false, RetryAfter: 0s, HTTPStatusCode: 0, RawError: test")
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -438,11 +438,11 @@ func TestCreateVolume(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &compute.DiskProperties{},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				rerr := &retry.Error{
 					RawError: fmt.Errorf(NotFound),
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(rerr).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(rerr).AnyTimes()
 				_, err := d.CreateVolume(context.Background(), req)
 				expectedErr := status.Error(codes.NotFound, "Retriable: false, RetryAfter: 0s, HTTPStatusCode: 0, RawError: NotFound")
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -474,8 +474,8 @@ func TestCreateVolume(t *testing.T) {
 						ProvisioningState: &state,
 					},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 				_, err := d.CreateVolume(context.Background(), req)
 				expectedErr := error(nil)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -533,8 +533,8 @@ func TestDeleteVolume(t *testing.T) {
 			ID: &id,
 		}
 
-		d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Eq(ctx), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
-		d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Delete(gomock.Eq(ctx), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Eq(ctx), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+		d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Delete(gomock.Eq(ctx), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 
 		result, err := d.DeleteVolume(context.Background(), test.req)
 		if err != nil {
@@ -632,7 +632,7 @@ func TestGetSnapshotInfo(t *testing.T) {
 
 func TestControllerPublishVolume(t *testing.T) {
 	d, err := NewFakeDriver(t)
-	d.cloud = &azure.Cloud{}
+	d.setCloud(&azure.Cloud{})
 	if err != nil {
 		t.Fatalf("Error getting driver: %v", err)
 	}
@@ -689,8 +689,8 @@ func TestControllerPublishVolume(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		defer ctrl.Finish()
 		mockDiskClient := mockdiskclient.NewMockInterface(ctrl)
-		d.cloud = &azure.Cloud{}
-		d.cloud.DisksClient = mockDiskClient
+		d.setCloud(&azure.Cloud{})
+		d.getCloud().DisksClient = mockDiskClient
 		mockDiskClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 		_, err := d.ControllerPublishVolume(context.Background(), test.req)
 		if !reflect.DeepEqual(err, test.expectedErr) {
@@ -701,7 +701,7 @@ func TestControllerPublishVolume(t *testing.T) {
 
 func TestControllerUnpublishVolume(t *testing.T) {
 	d, err := NewFakeDriver(t)
-	d.cloud = &azure.Cloud{}
+	d.setCloud(&azure.Cloud{})
 	if err != nil {
 		t.Fatalf("Error getting driver: %v", err)
 	}
@@ -749,7 +749,7 @@ func TestControllerGetCapabilities(t *testing.T) {
 	capList := []*csi.ControllerServiceCapability{{
 		Type: capType,
 	}}
-	d.Cap = capList
+	d.setControllerCapabilities(capList)
 	// Test valid request
 	req := csi.ControllerGetCapabilitiesRequest{}
 	resp, err := d.ControllerGetCapabilities(context.Background(), &req)
@@ -813,7 +813,7 @@ func TestControllerExpandVolume(t *testing.T) {
 				ctx := context.Background()
 				d, _ := NewFakeDriver(t)
 				var csc []*csi.ControllerServiceCapability
-				d.Cap = csc
+				d.setControllerCapabilities(csc)
 				expectedErr := status.Error(codes.InvalidArgument, "invalid expand volume request: volume_id:\"vol_1\" ")
 				_, err := d.ControllerExpandVolume(ctx, req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -891,8 +891,8 @@ func TestControllerExpandVolume(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockDiskClient := mockdiskclient.NewMockInterface(ctrl)
-				d.cloud = &azure.Cloud{}
-				d.cloud.DisksClient = mockDiskClient
+				d.setCloud(&azure.Cloud{})
+				d.getCloud().DisksClient = mockDiskClient
 				mockDiskClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				expectedErr := status.Errorf(codes.Internal, "could not get size of the disk(unit-test-volume)")
 				_, err := d.ControllerExpandVolume(ctx, req)
@@ -975,7 +975,7 @@ func TestCreateSnapshot(t *testing.T) {
 			name: "Invalid tag ",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				parameter := make(map[string]string)
 				parameter["tags"] = "unit-test"
 				parameter[incrementalField] = "false"
@@ -1001,11 +1001,11 @@ func TestCreateSnapshot(t *testing.T) {
 					Name:           "snapname",
 				}
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("test"),
 				}
@@ -1029,11 +1029,11 @@ func TestCreateSnapshot(t *testing.T) {
 					Parameters:     parameter,
 				}
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("existing disk"),
 				}
@@ -1056,11 +1056,11 @@ func TestCreateSnapshot(t *testing.T) {
 					Parameters:     parameter,
 				}
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("get snapshot error"),
 				}
@@ -1085,11 +1085,11 @@ func TestCreateSnapshot(t *testing.T) {
 					Parameters:     parameter,
 				}
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				provisioningState := "succeeded"
 				DiskSize := int32(10)
 				snapshotID := "test"
@@ -1164,11 +1164,11 @@ func TestDeleteSnapshot(t *testing.T) {
 			name: "delete Snapshot error",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				req := &csi.DeleteSnapshotRequest{
 					SnapshotId: "testurl/subscriptions/12/resourceGroups/23/providers/Microsoft.Compute/snapshots/snapshot-name",
 				}
@@ -1187,11 +1187,11 @@ func TestDeleteSnapshot(t *testing.T) {
 			name: "Valid delete Snapshot ",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				req := &csi.DeleteSnapshotRequest{
 					SnapshotId: "testurl/subscriptions/12/resourceGroups/23/providers/Microsoft.Compute/snapshots/snapshot-name",
 				}
@@ -1217,7 +1217,7 @@ func TestGenerateCSISnapshot(t *testing.T) {
 			name: "snap shot property not exist",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				snapshot := compute.Snapshot{}
 				sourceVolumeID := "unit-test"
 				_, err := generateCSISnapshot(sourceVolumeID, &snapshot)
@@ -1231,7 +1231,7 @@ func TestGenerateCSISnapshot(t *testing.T) {
 			name: "diskSizeGB of snapshot property is nil",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				provisioningState := "true"
 				snapshot := compute.Snapshot{
 					SnapshotProperties: &compute.SnapshotProperties{
@@ -1297,10 +1297,10 @@ func TestGetSnapshotByID(t *testing.T) {
 				sourceVolumeID := "unit-test"
 				ctx := context.Background()
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				snapshotID := "testurl/subscriptions/23/providers/Microsoft.Compute/snapshots/snapshot-name"
 				expectedErr := fmt.Errorf("could not get snapshot name from testurl/subscriptions/23/providers/Microsoft.Compute/snapshots/snapshot-name, correct format: (?i).*/subscriptions/(?:.*)/resourceGroups/(?:.*)/providers/Microsoft.Compute/snapshots/(.+)")
-				_, err := d.getSnapshotByID(ctx, d.cloud.ResourceGroup, snapshotID, sourceVolumeID)
+				_, err := d.getSnapshotByID(ctx, d.getCloud().ResourceGroup, snapshotID, sourceVolumeID)
 				if !reflect.DeepEqual(err, expectedErr) {
 					t.Errorf("actualErr: (%v), expectedErr: (%v)", err, expectedErr)
 				}
@@ -1310,11 +1310,11 @@ func TestGetSnapshotByID(t *testing.T) {
 			name: "snapshot get error",
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
-				d.cloud = &azure.Cloud{}
+				d.setCloud(&azure.Cloud{})
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("test"),
 				}
@@ -1326,7 +1326,7 @@ func TestGetSnapshotByID(t *testing.T) {
 				snapshotVolumeID := "unit-test"
 				mockSnapshotClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(snapshot, rerr).AnyTimes()
 				expectedErr := fmt.Errorf("could not get snapshot name from testurl/subscriptions/23/providers/Microsoft.Compute/snapshots/snapshot-name, correct format: (?i).*/subscriptions/(?:.*)/resourceGroups/(?:.*)/providers/Microsoft.Compute/snapshots/(.+)")
-				_, err := d.getSnapshotByID(context.Background(), d.cloud.ResourceGroup, snapshotID, snapshotVolumeID)
+				_, err := d.getSnapshotByID(context.Background(), d.getCloud().ResourceGroup, snapshotID, snapshotVolumeID)
 				if !reflect.DeepEqual(err, expectedErr) {
 					t.Errorf("actualErr: (%v), expectedErr: (%v)", err, expectedErr)
 				}
@@ -1378,7 +1378,7 @@ func TestListSnapshots(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				mockSnapshotClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(snapshot, nil).AnyTimes()
 				expectedErr := error(nil)
 				_, err := d.ListSnapshots(context.TODO(), &req)
@@ -1401,7 +1401,7 @@ func TestListSnapshots(t *testing.T) {
 					RawError: fmt.Errorf("test"),
 				}
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				mockSnapshotClient.EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(snapshots, rerr).AnyTimes()
 				expectedErr := status.Error(codes.Internal, "Unknown list snapshot error: Retriable: false, RetryAfter: 0s, HTTPStatusCode: 0, RawError: test")
 				_, err := d.ListSnapshots(context.TODO(), &req)
@@ -1421,7 +1421,7 @@ func TestListSnapshots(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				mockSnapshotClient.EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(snapshots, nil).AnyTimes()
 				expectedErr := fmt.Errorf("failed to generate snapshot entry: snapshot property is nil")
 				_, err := d.ListSnapshots(context.TODO(), &req)
@@ -1453,7 +1453,7 @@ func TestListSnapshots(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				defer ctrl.Finish()
 				mockSnapshotClient := mocksnapshotclient.NewMockInterface(ctrl)
-				d.cloud.SnapshotsClient = mockSnapshotClient
+				d.getCloud().SnapshotsClient = mockSnapshotClient
 				mockSnapshotClient.EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(snapshots, nil).AnyTimes()
 				snapshotsResponse, _ := d.ListSnapshots(context.TODO(), &req)
 				if len(snapshotsResponse.Entries) != 1 {
@@ -1519,7 +1519,7 @@ func TestListVolumes(t *testing.T) {
 				disk := compute.Disk{ID: &fakeVolumeID}
 				disks := []compute.Disk{}
 				disks = append(disks, disk)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1541,7 +1541,7 @@ func TestListVolumes(t *testing.T) {
 				disk1, disk2 := compute.Disk{ID: &fakeVolumeID}, compute.Disk{ID: &fakeVolumeID}
 				disks := []compute.Disk{}
 				disks = append(disks, disk1, disk2)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1567,7 +1567,7 @@ func TestListVolumes(t *testing.T) {
 				disk1, disk2 := compute.Disk{ID: &fakeVolumeID1}, compute.Disk{ID: &fakeVolumeID12}
 				disks := []compute.Disk{}
 				disks = append(disks, disk1, disk2)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1592,7 +1592,7 @@ func TestListVolumes(t *testing.T) {
 				}
 				d, _ := NewFakeDriver(t)
 				disks := []compute.Disk{}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, nil).AnyTimes()
 				expectedErr := status.Error(codes.FailedPrecondition, "ListVolumes starting token(1) on rg(rg) is greater than total number of volumes")
 				_, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1611,7 +1611,7 @@ func TestListVolumes(t *testing.T) {
 				rerr := &retry.Error{
 					RawError: fmt.Errorf("test"),
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, rerr).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return(disks, rerr).AnyTimes()
 				expectedErr := status.Error(codes.Internal, "ListVolumes on rg(rg) failed with error: Retriable: false, RetryAfter: 0s, HTTPStatusCode: 0, RawError: test")
 				_, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1628,9 +1628,9 @@ func TestListVolumes(t *testing.T) {
 				pvList := v1.PersistentVolumeList{
 					Items: []v1.PersistentVolume{volume1},
 				}
-				d.cloud.KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
+				d.getCloud().KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
 				disk1 := compute.Disk{ID: &fakeVolumeID}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1653,9 +1653,9 @@ func TestListVolumes(t *testing.T) {
 				pvList := v1.PersistentVolumeList{
 					Items: []v1.PersistentVolume{volume1, volume2},
 				}
-				d.cloud.KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk2}, nil)
+				d.getCloud().KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk2}, nil)
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1682,9 +1682,9 @@ func TestListVolumes(t *testing.T) {
 				}
 				fakeVolumeID11, fakeVolumeID12 := "/subscriptions/test-subscription/resourceGroups/test_resourcegroup-1/providers/Microsoft.Compute/disks/test-pv-1", "/subscriptions/test-subscription/resourceGroups/test_resourcegroup-2/providers/Microsoft.Compute/disks/test-pv-2"
 				disk1, disk2 := compute.Disk{ID: &fakeVolumeID11}, compute.Disk{ID: &fakeVolumeID12}
-				d.cloud.KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk2}, nil)
+				d.getCloud().KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk1}, nil)
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().ListByResourceGroup(gomock.Any(), gomock.Any()).Return([]compute.Disk{disk2}, nil)
 				expectedErr := error(nil)
 				listVolumesResponse, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1711,7 +1711,7 @@ func TestListVolumes(t *testing.T) {
 				pvList := v1.PersistentVolumeList{
 					Items: []v1.PersistentVolume{},
 				}
-				d.cloud.KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
+				d.getCloud().KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(&pvList, nil)
 				expectedErr := status.Error(codes.FailedPrecondition, "ListVolumes starting token(1) is greater than total number of disks")
 				_, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1727,7 +1727,7 @@ func TestListVolumes(t *testing.T) {
 				}
 				d := getFakeDriverWithKubeClient(t)
 				rerr := fmt.Errorf("test")
-				d.cloud.KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(nil, rerr)
+				d.getCloud().KubeClient.CoreV1().PersistentVolumes().(*mockpersistentvolume.MockInterface).EXPECT().List(gomock.Any(), gomock.Any()).Return(nil, rerr)
 				expectedErr := status.Error(codes.Internal, "ListVolumes failed while fetching PersistentVolumes List with error: test")
 				_, err := d.ListVolumes(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1799,7 +1799,7 @@ func TestValidateVolumeCapabilities(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &compute.DiskProperties{},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				expectedErr := error(nil)
 				_, err := d.ValidateVolumeCapabilities(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1826,7 +1826,7 @@ func TestValidateVolumeCapabilities(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &compute.DiskProperties{},
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				expectedErr := error(nil)
 				_, err := d.ValidateVolumeCapabilities(context.TODO(), &req)
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1937,7 +1937,7 @@ func TestGetSourceDiskSize(t *testing.T) {
 			testFunc: func(t *testing.T) {
 				d, _ := NewFakeDriver(t)
 				disk := compute.Disk{}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				_, err := d.GetSourceDiskSize(context.Background(), "test-rg", "test-disk", 0, 1)
 				expectedErr := status.Error(codes.Internal, "DiskProperty not found for disk (test-disk) in resource group (test-rg)")
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1953,7 +1953,7 @@ func TestGetSourceDiskSize(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &diskProperties,
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				_, err := d.GetSourceDiskSize(context.Background(), "test-rg", "test-disk", 0, 1)
 				expectedErr := status.Error(codes.Internal, "DiskSizeGB for disk (test-disk) in resourcegroup (test-rg) is nil")
 				if !reflect.DeepEqual(err, expectedErr) {
@@ -1972,7 +1972,7 @@ func TestGetSourceDiskSize(t *testing.T) {
 				disk := compute.Disk{
 					DiskProperties: &diskProperties,
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk, nil).AnyTimes()
 				size, _ := d.GetSourceDiskSize(context.Background(), "test-rg", "test-disk", 0, 1)
 				expectedOutput := diskSizeGB
 				if *size != expectedOutput {
@@ -2004,7 +2004,7 @@ func TestGetSourceDiskSize(t *testing.T) {
 				disk2 := compute.Disk{
 					DiskProperties: &diskProperties2,
 				}
-				d.cloud.DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk1, nil).Return(disk2, nil).AnyTimes()
+				d.getCloud().DisksClient.(*mockdiskclient.MockInterface).EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any()).Return(disk1, nil).Return(disk2, nil).AnyTimes()
 				size, _ := d.GetSourceDiskSize(context.Background(), "test-rg", "test-disk-1", 0, 2)
 				expectedOutput := diskSizeGB2
 				if *size != expectedOutput {
@@ -2018,14 +2018,14 @@ func TestGetSourceDiskSize(t *testing.T) {
 	}
 }
 
-func getFakeDriverWithKubeClient(t *testing.T) *Driver {
+func getFakeDriverWithKubeClient(t *testing.T) FakeDriver {
 	d, _ := NewFakeDriver(t)
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 	corev1 := mockcorev1.NewMockInterface(ctrl)
 	persistentvolume := mockpersistentvolume.NewMockInterface(ctrl)
-	d.cloud.KubeClient = mockkubeclient.NewMockInterface(ctrl)
-	d.cloud.KubeClient.(*mockkubeclient.MockInterface).EXPECT().CoreV1().Return(corev1).AnyTimes()
-	d.cloud.KubeClient.CoreV1().(*mockcorev1.MockInterface).EXPECT().PersistentVolumes().Return(persistentvolume).AnyTimes()
+	d.getCloud().KubeClient = mockkubeclient.NewMockInterface(ctrl)
+	d.getCloud().KubeClient.(*mockkubeclient.MockInterface).EXPECT().CoreV1().Return(corev1).AnyTimes()
+	d.getCloud().KubeClient.CoreV1().(*mockcorev1.MockInterface).EXPECT().PersistentVolumes().Return(persistentvolume).AnyTimes()
 	return d
 }
diff --git a/pkg/azuredisk/fake_mounter_test.go b/pkg/azuredisk/fake_mounter_test.go
index d38d91b0b..c35d48525 100644
--- a/pkg/azuredisk/fake_mounter_test.go
+++ b/pkg/azuredisk/fake_mounter_test.go
@@ -53,11 +53,11 @@ func TestMount(t *testing.T) {
 
 	d, _ := NewFakeDriver(t)
 	fakeMounter := &fakeMounter{}
-	d.mounter = &mount.SafeFormatAndMount{
+	d.setMounter(&mount.SafeFormatAndMount{
 		Interface: fakeMounter,
-	}
+	})
 	for _, test := range tests {
-		err := d.mounter.Mount(test.source, test.target, "", nil)
+		err := d.getMounter().Mount(test.source, test.target, "", nil)
 		if !reflect.DeepEqual(err, test.expectedErr) {
 			t.Errorf("Unexpected error: %v", err)
 		}
@@ -93,11 +93,11 @@ func TestMountSensitive(t *testing.T) {
 
 	d, _ := NewFakeDriver(t)
 	fakeMounter := &fakeMounter{}
-	d.mounter = &mount.SafeFormatAndMount{
+	d.setMounter(&mount.SafeFormatAndMount{
 		Interface: fakeMounter,
-	}
+	})
 	for _, test := range tests {
-		err := d.mounter.MountSensitive(test.source, test.target, "", nil, nil)
+		err := d.getMounter().MountSensitive(test.source, test.target, "", nil, nil)
 		if !reflect.DeepEqual(err, test.expectedErr) {
 			t.Errorf("Unexpected error: %v", err)
 		}
@@ -128,11 +128,11 @@ func TestIsLikelyNotMountPoint(t *testing.T) {
 
 	d, _ := NewFakeDriver(t)
 	fakeMounter := &fakeMounter{}
-	d.mounter = &mount.SafeFormatAndMount{
+	d.setMounter(&mount.SafeFormatAndMount{
 		Interface: fakeMounter,
-	}
+	})
 	for _, test := range tests {
-		_, err := d.mounter.IsLikelyNotMountPoint(test.file)
+		_, err := d.getMounter().IsLikelyNotMountPoint(test.file)
 		if !reflect.DeepEqual(err, test.expectedErr) {
 			t.Errorf("Unexpected error: %v", err)
 		}
diff --git a/pkg/azuredisk/identityserver_test.go b/pkg/azuredisk/identityserver_test.go
index f4a77d6a1..e01a242c4 100644
--- a/pkg/azuredisk/identityserver_test.go
+++ b/pkg/azuredisk/identityserver_test.go
@@ -18,9 +18,10 @@ package azuredisk
 
 import (
 	"context"
+	"testing"
+
 	"github.com/container-storage-interface/spec/lib/go/csi"
 	"github.com/stretchr/testify/assert"
-	"testing"
 )
 
 const (
@@ -39,7 +40,7 @@ func TestGetPluginInfo(t *testing.T) {
 
 	//Check error when driver name is empty
 	d, _ = NewFakeDriver(t)
-	d.Name = ""
+	d.setName("")
 	req = csi.GetPluginInfoRequest{}
 	resp, err = d.GetPluginInfo(context.Background(), &req)
 	assert.Error(t, err)
@@ -47,7 +48,7 @@ func TestGetPluginInfo(t *testing.T) {
 
 	//Check error when version is empty
 	d, _ = NewFakeDriver(t)
-	d.Version = ""
+	d.setVersion("")
 	req = csi.GetPluginInfoRequest{}
 	resp, err = d.GetPluginInfo(context.Background(), &req)
 	assert.Error(t, err)
diff --git a/pkg/azuredisk/nodeserver_test.go b/pkg/azuredisk/nodeserver_test.go
index 89482d60f..a4a70e3b7 100644
--- a/pkg/azuredisk/nodeserver_test.go
+++ b/pkg/azuredisk/nodeserver_test.go
@@ -34,7 +34,6 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
-	"sigs.k8s.io/azuredisk-csi-driver/pkg/mounter"
 	volumehelper "sigs.k8s.io/azuredisk-csi-driver/pkg/util"
 )
 
@@ -70,7 +69,7 @@ func TestNodeGetCapabilities(t *testing.T) {
 	capList := []*csi.NodeServiceCapability{{
 		Type: capType,
 	}}
-	d.NSCap = capList
+	d.setNodeCapabilities(capList)
 	// Test valid request
 	req := csi.NodeGetCapabilitiesRequest{}
 	resp, err := d.NodeGetCapabilities(context.Background(), &req)
@@ -199,7 +198,7 @@ func TestEnsureMountPoint(t *testing.T) {
 	d, _ := NewFakeDriver(t)
 	fakeMounter, err := NewFakeMounter()
 	assert.NoError(t, err)
-	d.mounter = fakeMounter
+	d.setMounter(fakeMounter)
 
 	for _, test := range tests {
 		if !(runtime.GOOS == "windows" && test.skipOnWindows) {
@@ -344,7 +343,10 @@ func TestNodeStageVolume(t *testing.T) {
 	_ = makeDir(sourceTest)
 	_ = makeDir(targetTest)
 	d, _ := NewFakeDriver(t)
-	d.mounter, _ = mounter.NewSafeMounter()
+	fakeMounter, err := NewFakeMounter()
+	assert.NoError(t, err)
+	d.setMounter(fakeMounter)
+
 	for _, test := range tests {
 		if !(test.skipOnDarwin && runtime.GOOS == "darwin") {
 			_, err := d.NodeStageVolume(context.Background(), &test.req)
@@ -357,7 +359,7 @@ func TestNodeStageVolume(t *testing.T) {
 	}
 
 	// Clean up
-	err := os.RemoveAll(sourceTest)
+	err = os.RemoveAll(sourceTest)
 	assert.NoError(t, err)
 	err = os.RemoveAll(targetTest)
 	assert.NoError(t, err)
@@ -412,7 +414,7 @@ func TestNodeUnstageVolume(t *testing.T) {
 	d, _ := NewFakeDriver(t)
 	fakeMounter, err := NewFakeMounter()
 	assert.NoError(t, err)
-	d.mounter = fakeMounter
+	d.setMounter(fakeMounter)
 
 	for _, test := range tests {
 		if !(runtime.GOOS == "windows" && test.skipOnWindows) &&
@@ -569,7 +571,7 @@ func TestNodePublishVolume(t *testing.T) {
 	assert.NoError(t, err)
 	fakeMounter, err := NewFakeMounter()
 	assert.NoError(t, err)
-	d.mounter = fakeMounter
+	d.setMounter(fakeMounter)
 
 	for _, test := range tests {
 		if !(test.skipOnWindows && runtime.GOOS == "windows") {
@@ -635,7 +637,7 @@ func TestNodeUnpublishVolume(t *testing.T) {
 	d, _ := NewFakeDriver(t)
 	fakeMounter, err := NewFakeMounter()
 	assert.NoError(t, err)
-	d.mounter = fakeMounter
+	d.setMounter(fakeMounter)
 
 	for _, test := range tests {
 		if !(test.skipOnWindows && runtime.GOOS == "windows") &&
@@ -799,7 +801,7 @@ func TestGetBlockSizeBytes(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		_, err := getBlockSizeBytes(test.req, d.mounter)
+		_, err := getBlockSizeBytes(test.req, d.getMounter())
 		if !testutil.AssertError(&test.expectedErr, err) {
 			t.Errorf("desc: %s\n actualErr: (%v), expectedErr: (%v)", test.desc, err, test.expectedErr.Error())
 		}
@@ -934,7 +936,7 @@ func TestGetDevicePathWithMountPath(t *testing.T) {
 
 	for _, test := range tests {
 		if !(test.skipOnDarwin && runtime.GOOS == "darwin") && !(test.skipOnWindows && runtime.GOOS == "windows") {
-			_, err := getDevicePathWithMountPath(test.req, d.mounter)
+			_, err := getDevicePathWithMountPath(test.req, d.getMounter())
 			if !reflect.DeepEqual(err, test.expectedErr) {
 				t.Errorf("desc: %s\n actualErr: (%v), expectedErr: (%v)", test.desc, err, test.expectedErr)
 			}
diff --git a/test/e2e/suite_test.go b/test/e2e/suite_test.go
index 555ccf88a..ef3206942 100644
--- a/test/e2e/suite_test.go
+++ b/test/e2e/suite_test.go
@@ -51,7 +51,7 @@ const (
 )
 
 var (
-	azurediskDriver           *azuredisk.Driver
+	azurediskDriver           azuredisk.CSIDriver
 	isUsingInTreeVolumePlugin = os.Getenv(driver.AzureDriverNameVar) == inTreeStorageClass
 	isTestingMigration        = os.Getenv(testMigrationEnvVar) != ""
 	isWindowsCluster          = os.Getenv(testWindowsEnvVar) != ""
diff --git a/test/e2e/testsuites/dynamically_provisioned_reclaim_policy_tester.go b/test/e2e/testsuites/dynamically_provisioned_reclaim_policy_tester.go
index b5818f6fc..3a8588dfa 100644
--- a/test/e2e/testsuites/dynamically_provisioned_reclaim_policy_tester.go
+++ b/test/e2e/testsuites/dynamically_provisioned_reclaim_policy_tester.go
@@ -30,7 +30,7 @@ import (
 type DynamicallyProvisionedReclaimPolicyTest struct {
 	CSIDriver              driver.DynamicPVTestDriver
 	Volumes                []VolumeDetails
-	Azuredisk              *azuredisk.Driver
+	Azuredisk              azuredisk.CSIDriver
 	StorageClassParameters map[string]string
 }
 
diff --git a/test/e2e/testsuites/pre_provisioned_dangling_attach_volume_tester.go b/test/e2e/testsuites/pre_provisioned_dangling_attach_volume_tester.go
index e880ac79c..769e39df1 100644
--- a/test/e2e/testsuites/pre_provisioned_dangling_attach_volume_tester.go
+++ b/test/e2e/testsuites/pre_provisioned_dangling_attach_volume_tester.go
@@ -33,7 +33,7 @@ import (
 // Testing that a volume could be reattached to a different node on DanglingAttachError
 type PreProvisionedDanglingAttachVolumeTest struct {
 	CSIDriver       driver.PreProvisionedVolumeTestDriver
-	AzureDiskDriver *azuredisk.Driver
+	AzureDiskDriver azuredisk.CSIDriver
 	Pod             PodDetails
 	VolumeContext   map[string]string
 }
diff --git a/test/e2e/testsuites/testsuites.go b/test/e2e/testsuites/testsuites.go
index c13343a65..181dced9d 100644
--- a/test/e2e/testsuites/testsuites.go
+++ b/test/e2e/testsuites/testsuites.go
@@ -384,7 +384,7 @@ func (t *TestPersistentVolumeClaim) DeleteBoundPersistentVolume() {
 	framework.ExpectNoError(err)
 }
 
-func (t *TestPersistentVolumeClaim) DeleteBackingVolume(driver *azuredisk.Driver) {
+func (t *TestPersistentVolumeClaim) DeleteBackingVolume(driver azuredisk.CSIDriver) {
 	volumeID := t.persistentVolume.Spec.CSI.VolumeHandle
 	ginkgo.By(fmt.Sprintf("deleting azuredisk volume %q", volumeID))
 	req := &csi.DeleteVolumeRequest{
diff --git a/test/integration/integration_test.go b/test/integration/integration_test.go
index c7351de9c..92423f5eb 100644
--- a/test/integration/integration_test.go
+++ b/test/integration/integration_test.go
@@ -18,6 +18,7 @@ package integration
 
 import (
 	"context"
+	"flag"
 	"log"
 	"os"
 	"os/exec"
@@ -34,6 +35,8 @@ const (
 	nodeid = "integration-test-node"
 )
 
+var useDriverV2 = flag.Bool("temp-use-driver-v2", false, "A temporary flag to enable early test and development of Azure Disk CSI Driver V2. This will be removed in the future.")
+
 func TestIntegrationOnAzurePublicCloud(t *testing.T) {
 	// Test on AzurePublicCloud
 	creds, err := credentials.CreateAzureCredentialFile()
@@ -82,7 +85,12 @@ func TestIntegrationOnAzurePublicCloud(t *testing.T) {
 	assert.NoError(t, err)
 	assert.True(t, strings.HasSuffix(cwd, "azuredisk-csi-driver"))
 
-	cmd := exec.Command("./test/integration/run-tests-all-clouds.sh", creds.Cloud)
+	args := []string{creds.Cloud}
+	if *useDriverV2 {
+		args = append(args, "v2")
+	}
+
+	cmd := exec.Command("./test/integration/run-tests-all-clouds.sh", args...)
 	cmd.Dir = cwd
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
diff --git a/test/integration/run-test.sh b/test/integration/run-test.sh
index 150a9cb48..8265e5ee3 100755
--- a/test/integration/run-test.sh
+++ b/test/integration/run-test.sh
@@ -42,7 +42,11 @@ fi
 echo "Begin to run integration test on $cloud..."
 
 # Run CSI driver as a background service
-_output/azurediskplugin --endpoint "$endpoint" --nodeid "$node" -v=5 &
+if [[ $# -lt 4 || "$4" != "v2" ]]; then
+  _output/azurediskplugin --endpoint "$endpoint" --nodeid "$node" -v=5 &
+else
+  _output/azurediskpluginv2 --endpoint "$endpoint" --nodeid "$node" -v=5 --temp-use-driver-v2 &
+fi
 trap cleanup EXIT
 
 if [[ "$cloud" == 'AzureChinaCloud' ]]; then
diff --git a/test/integration/run-tests-all-clouds.sh b/test/integration/run-tests-all-clouds.sh
index b7374d980..4d8e794d1 100755
--- a/test/integration/run-tests-all-clouds.sh
+++ b/test/integration/run-tests-all-clouds.sh
@@ -24,4 +24,9 @@ if [[ "$#" -gt 0 ]]; then
   cloud="$1"
 fi
 
-test/integration/run-test.sh 'tcp://127.0.0.1:10000' "$nodeid" "$cloud"
+version='v1'
+if [[ "$#" -gt 1 ]]; then
+  version="$2"
+fi
+
+test/integration/run-test.sh 'tcp://127.0.0.1:10000' "$nodeid" "$cloud" "$version"
diff --git a/test/sanity/run-test.sh b/test/sanity/run-test.sh
index b010ffe35..3e15fb667 100755
--- a/test/sanity/run-test.sh
+++ b/test/sanity/run-test.sh
@@ -31,7 +31,11 @@ if [[ "$#" -gt 0 ]] && [[ -n "$1" ]]; then
   nodeid="$1"
 fi
 
-_output/azurediskplugin --endpoint "$endpoint" --nodeid "$nodeid" -v=5 &
+if [[ "$#" -lt 2 || "$2" != "v2" ]]; then
+  _output/azurediskplugin --endpoint "$endpoint" --nodeid "$nodeid" -v=5 &
+else
+  _output/azurediskpluginv2 --endpoint "$endpoint" --nodeid "$nodeid" -v=5 --temp-use-driver-v2 &
+fi
 
 echo 'Begin to run sanity test...'
 readonly CSI_SANITY_BIN='csi-sanity'
diff --git a/test/sanity/run-tests-all-clouds.sh b/test/sanity/run-tests-all-clouds.sh
index d4c11f407..aa9478720 100755
--- a/test/sanity/run-tests-all-clouds.sh
+++ b/test/sanity/run-tests-all-clouds.sh
@@ -31,4 +31,4 @@ function install_csi_sanity_bin {
 if [[ -z "$(command -v csi-sanity)" ]]; then
 	install_csi_sanity_bin
 fi
-test/sanity/run-test.sh "$nodeid"
+test/sanity/run-test.sh "$nodeid" $*
diff --git a/test/sanity/sanity_test.go b/test/sanity/sanity_test.go
index 78149dfea..3b01bcef0 100644
--- a/test/sanity/sanity_test.go
+++ b/test/sanity/sanity_test.go
@@ -18,6 +18,7 @@ package sanity
 
 import (
 	"context"
+	"flag"
 	"log"
 	"os"
 	"os/exec"
@@ -35,6 +36,8 @@ const (
 	vmType = "standard"
 )
 
+var useDriverV2 = flag.Bool("temp-use-driver-v2", false, "A temporary flag to enable early test and development of Azure Disk CSI Driver V2. This will be removed in the future.")
+
 func TestSanity(t *testing.T) {
 	// Set necessary env vars for creating azure credential file
 	os.Setenv("AZURE_VM_TYPE", vmType)
@@ -85,7 +88,12 @@ func TestSanity(t *testing.T) {
 	assert.NoError(t, err)
 	assert.True(t, strings.HasSuffix(projectRoot, "azuredisk-csi-driver"))
 
-	cmd := exec.Command("./test/sanity/run-tests-all-clouds.sh")
+	args := make([]string, 0)
+	if *useDriverV2 {
+		args = append(args, "v2")
+	}
+
+	cmd := exec.Command("./test/sanity/run-tests-all-clouds.sh", args...)
 	cmd.Dir = projectRoot
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr