diff --git a/sdk/compute/Microsoft.Azure.Management.Compute/tests/DiskRPTests/DiskRPTestsBase.cs b/sdk/compute/Microsoft.Azure.Management.Compute/tests/DiskRPTests/DiskRPTestsBase.cs
index 34e3b498757e0..05bf6fa5803c6 100644
--- a/sdk/compute/Microsoft.Azure.Management.Compute/tests/DiskRPTests/DiskRPTestsBase.cs
+++ b/sdk/compute/Microsoft.Azure.Management.Compute/tests/DiskRPTests/DiskRPTestsBase.cs
@@ -1066,7 +1066,7 @@ protected DiskEncryptionSet GenerateDefaultDiskEncryptionSet(string location, st
                     Type = ResourceIdentityType.SystemAssigned.ToString()
                 },
                 Location = location,
-                ActiveKey = new KeyVaultAndKeyReference
+                ActiveKey = new KeyForDiskEncryptionSet
                 {
                     SourceVault = new SourceVault
                     {