diff --git a/src/ResourceManager/Compute/Commands.Compute/Extension/Diagnostics/SetAzureVMDiagnosticsExtension.cs b/src/ResourceManager/Compute/Commands.Compute/Extension/Diagnostics/SetAzureVMDiagnosticsExtension.cs
index ba28ab273139..0864f5feafcc 100644
--- a/src/ResourceManager/Compute/Commands.Compute/Extension/Diagnostics/SetAzureVMDiagnosticsExtension.cs
+++ b/src/ResourceManager/Compute/Commands.Compute/Extension/Diagnostics/SetAzureVMDiagnosticsExtension.cs
@@ -21,7 +21,8 @@
 using Microsoft.Azure.Management.Compute.Models;
 using Microsoft.WindowsAzure.Commands.Common.Storage;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
-using Microsoft.WindowsAzure.Management.Storage;
+using Microsoft.Azure.Commands.Management.Storage;
+using Microsoft.Azure.Management.Storage;
 using Newtonsoft.Json;
 
 namespace Microsoft.Azure.Commands.Compute
@@ -37,7 +38,7 @@ public class SetAzureVMDiagnosticsExtensionCommand : VirtualMachineExtensionBase
         private const string VirtualMachineExtension = "Microsoft.Compute/virtualMachines/extensions";
         private const string IaaSDiagnosticsExtension = "IaaSDiagnostics";
         private const string ExtensionPublisher = "Microsoft.Azure.Diagnostics";
-        private StorageManagementClient storageClient;
+        private IStorageManagementClient storageClient;
 
         [Parameter(
             Mandatory = true,
@@ -154,14 +155,14 @@ public string PrivateConfiguration
             }
         }
 
-        public StorageManagementClient StorageClient
+        public IStorageManagementClient StorageClient
         {
             get
             {
                 if (this.storageClient == null)
                 {
                     this.storageClient = AzureSession.ClientFactory.CreateClient<StorageManagementClient>(
-                        DefaultProfile.Context, AzureEnvironment.Endpoint.ServiceManagement);
+                        DefaultProfile.Context, AzureEnvironment.Endpoint.ResourceManager);
                 }
 
                 return this.storageClient;
@@ -201,15 +202,8 @@ protected string GetStorageKey()
 
             if (!string.IsNullOrEmpty(StorageAccountName))
             {
-                var storageAccount = this.StorageClient.StorageAccounts.Get(StorageAccountName);
-                if (storageAccount != null)
-                {
-                    var keys = this.StorageClient.StorageAccounts.GetKeys(StorageAccountName);
-                    if (keys != null)
-                    {
-                        storageKey = !string.IsNullOrEmpty(keys.PrimaryKey) ? keys.PrimaryKey : keys.SecondaryKey;
-                    }
-                }
+                var storageCredentials = StorageUtilities.GenerateStorageCredentials(this.StorageClient, this.ResourceGroupName, this.StorageAccountName);
+                storageKey = storageCredentials.ExportBase64EncodedKey();
             }
 
             return storageKey;
diff --git a/tools/Docs/Get-AllCommands/Get-AllCommands.ps1 b/tools/Docs/Get-AllCommands/Get-AllCommands.ps1
index f381c7d77a4f..357d1897dd01 100644
--- a/tools/Docs/Get-AllCommands/Get-AllCommands.ps1
+++ b/tools/Docs/Get-AllCommands/Get-AllCommands.ps1
@@ -690,9 +690,9 @@ function Get-AllBuildServerCommands {
 rm .\Output -Recurse -Force -ErrorAction SilentlyContinue
 
 #Comment these lines to selectivly build the output for either Service Management or Resource Manager
-Get-AllBuildServerCommands -OutputPath ".\Output" -ManifestFullName "..\..\..\src\Package\Debug\ServiceManagement\Azure\Azure.psd1"
+Get-AllBuildServerCommands -OutputPath ".\Output" -ManifestFullName "..\..\..\src\Package\Release\ServiceManagement\Azure\Azure.psd1"
 Get-AllBuildServerCommands -OutputPath ".\Output" -ManifestFullName "..\..\AzureRM\AzureRM.psd1"
 
-$modules = (Get-ChildItem "..\..\..\src\Package\Debug\ResourceManager" -Recurse -Include "*.psd1" -Exclude "*dll-help.psd1", "AzureResourceManager.psd1") | sort -Unique -Property Name
+$modules = (Get-ChildItem "..\..\..\src\Package\Release\ResourceManager" -Recurse -Include "*.psd1" -Exclude "*dll-help.psd1", "AzureResourceManager.psd1") | sort -Unique -Property Name
 $modules | Foreach { Get-AllBuildServerCommands -OutputPath ".\Output" -ManifestFullName $_.FullName }