Ensure that TokenCache is transmitted to context when context is set

Author: markcowl

src/ResourceManager/Common/Commands.ResourceManager.Common/AzureRMProfileExtensions.cs

@@ -17,7 +17,9 @@
 using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
+using Microsoft.Azure.Commands.ResourceManager.Common.Properties;
 using Microsoft.Azure.Common.Authentication.Models;
+using Microsoft.IdentityModel.Clients.ActiveDirectory;
 
 namespace Microsoft.Azure.Commands.ResourceManager.Common
 {
@@ -27,15 +29,20 @@ public static class AzureRMProfileExtensions
         /// Set the context for the current profile, preserving token cache information
         /// </summary>
         /// <param name="profile">The profile to change the context for</param>
-        /// <param name="newContext">The new context</param>
-        public static void SetContext(this AzureRMProfile profile, AzureContext newContext)
+        /// <param name="newContext">The new context, with no token cache information.</param>
+        public static void SetContextWithCache(this AzureRMProfile profile, AzureContext newContext)
         {
-            var currentContext = profile.Context;
-            if (currentContext != null && currentContext.TokenCache != null && currentContext.TokenCache.Length > 0)
+            if (profile == null)
             {
-                newContext.TokenCache = currentContext.TokenCache;
+                throw new ArgumentNullException("profile", Resources.ProfileCannotBeNull);
             }
 
+            if (newContext == null)
+            {
+                throw new ArgumentNullException("newContext", Resources.ContextCannotBeNull);
+            }
+
+            newContext.TokenCache = TokenCache.DefaultShared.Serialize();
             profile.Context = newContext;
         }
     }

src/ResourceManager/Common/Commands.ResourceManager.Common/Properties/Resources.Designer.cs

@@ -117,6 +117,9 @@
   <resheader name="writer">
     <value>System.Resources.ResXResourceWriter, System.Windows.Forms, Version=4.0.0.0, Culture=neutral, PublicKeyToken=b77a5c561934e089</value>
   </resheader>
+  <data name="ContextCannotBeNull" xml:space="preserve">
+    <value>Context cannot be null.  Please log in using Add-AzureRmAccount.</value>
+  </data>
   <data name="DataCollectionActivity" xml:space="preserve">
     <value>Microsoft Azure PowerShell Data Collection Confirmation</value>
   </data>
@@ -151,6 +154,9 @@ Select Y to enable data collection [Y/N]:</value>
   <data name="NoSubscriptionFound" xml:space="preserve">
     <value>The provided account {0} does not have access to any subscriptions. Please try logging in with different credentials.</value>
   </data>
+  <data name="ProfileCannotBeNull" xml:space="preserve">
+    <value>Profile cannot be null.  Please run Add-AzureRmAccount.</value>
+  </data>
   <data name="ResourceProviderRegisterAttempt" xml:space="preserve">
     <value>Attempting to register resource provider '{0}'</value>
   </data>

src/ResourceManager/Common/Commands.ResourceManager.Common/Properties/Resources.resx

@@ -19,6 +19,7 @@
 using System;
 using Microsoft.WindowsAzure.Commands.Common.Test.Mocks;
 using System.Collections.Generic;
+using Microsoft.IdentityModel.Clients.ActiveDirectory;
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 
 namespace Microsoft.Azure.Commands.ResourceManager.Common.Test
@@ -131,5 +132,18 @@ public void NoSubscriptionsInListThrows()
             AzureSubscription subValue;
             Assert.False(client.TryGetSubscription(DefaultTenant.ToString(), DefaultSubscription.ToString(), out subValue));
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+      public void SetContextPreservesTokenCache()
+        {
+            AzureRMProfile profile = null;
+            AzureContext context = new AzureContext(null, null, null, null);
+            Assert.Throws<ArgumentNullException>(() =>profile.SetContextWithCache(context));
+            profile = new AzureRMProfile();
+            Assert.Throws<ArgumentNullException>(() => profile.SetContextWithCache(null));
+            profile.SetContextWithCache(context);
+            Assert.Equal(TokenCache.DefaultShared.Serialize(), profile.Context.TokenCache);
+        }
     }
 }

src/ResourceManager/Profile/Commands.Profile.Test/AzureRMProfileTests.cs

@@ -55,7 +55,7 @@ protected override void ProcessRecord()
         {
             if (ParameterSetName == ContextParameterSet)
             {
-                AzureRmProfileProvider.Instance.Profile.SetContext(new AzureContext(Context.Subscription, Context.Account,
+                AzureRmProfileProvider.Instance.Profile.SetContextWithCache(new AzureContext(Context.Subscription, Context.Account,
                     Context.Environment, Context.Tenant));
             }
             else

src/ResourceManager/Profile/Commands.Profile/Context/SetAzureRMContext.cs

@@ -105,7 +105,7 @@ public AzureContext SetCurrentContext(string subscriptionId, string tenantId)
         {
             if (!string.IsNullOrWhiteSpace(tenantId))
             {
-                _profile.SetContext(new AzureContext(
+                _profile.SetContextWithCache(new AzureContext(
                     _profile.Context.Subscription,
                     _profile.Context.Account,
                     _profile.Context.Environment,
@@ -139,7 +139,7 @@ public AzureContext SetCurrentContext(string subscriptionId, string tenantId)
                     newSubscription.Name = null;
                 }
 
-                _profile.SetContext(new AzureContext(
+                _profile.SetContextWithCache(new AzureContext(
                     newSubscription,
                     _profile.Context.Account,
                     _profile.Context.Environment,