Fix cloud connection token update logic (#217)

Azure · Aug 28, 2018 · 177b7c8 · 177b7c8
1 parent f5a4251
commit 177b7c8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 22 deletions.
diff --git a/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/CloudConnection.cs b/edge-hub/src/Microsoft.Azure.Devices.Edge.Hub.CloudProxy/CloudConnection.cs
@@ -277,13 +277,10 @@ async Task<string> GetNewToken(string iotHub, string id, string currentToken, II
                 }
                 else
                 {
-                    Events.TokenNotUsable(iotHub, id, token);
-                    if (retrying)
-                    {
-                        await Task.Delay(TokenRetryWaitTime);
-                    }
+                    Events.TokenNotUsable(iotHub, id, token);                    
                 }
 
+                bool newTokenGetterCreated = false;
                 // No need to lock here as the lock is being held by the refresher.
                 TaskCompletionSource<string> tcs = this.tokenGetter
                     .GetOrElse(
@@ -292,9 +289,19 @@ async Task<string> GetNewToken(string iotHub, string id, string currentToken, II
                             Events.SafeCreateNewToken(id);
                             var taskCompletionSource = new TaskCompletionSource<string>();
                             this.tokenGetter = Option.Some(taskCompletionSource);
-                            this.connectionStatusChangedHandler(currentIdentity.Id, CloudConnectionStatus.TokenNearExpiry);
+                            newTokenGetterCreated = true;                            
                             return taskCompletionSource;
                         });
+
+                if (newTokenGetterCreated)
+                {
+                    if (retrying)
+                    {
+                        await Task.Delay(TokenRetryWaitTime);
+                    }
+                    this.connectionStatusChangedHandler(currentIdentity.Id, CloudConnectionStatus.TokenNearExpiry);
+                }
+
                 retrying = true;
                 token = await tcs.Task;
             }

diff --git a/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/CloudConnectionTest.cs b/edge-hub/test/Microsoft.Azure.Devices.Edge.Hub.CloudProxy.Test/CloudConnectionTest.cs
@@ -287,9 +287,9 @@ public async Task UpdateDeviceConnectionTest()
             string hostname = "dummy.azure-devices.net";
             string deviceId = "device1";
 
-            IClientCredentials GetClientCredentials()
+            IClientCredentials GetClientCredentials(TimeSpan tokenExpiryDuration)
             {
-                string token = TokenHelper.CreateSasToken(hostname, DateTime.UtcNow.AddSeconds(10));
+                string token = TokenHelper.CreateSasToken(hostname, DateTime.UtcNow.AddSeconds(tokenExpiryDuration.TotalSeconds));
                 var identity = new DeviceIdentity(hostname, deviceId);
                 return new TokenCredentials(identity, token, string.Empty);
             }
@@ -338,7 +338,7 @@ IClient GetMockedDeviceClient()
             ICloudConnectionProvider cloudConnectionProvider = new CloudConnectionProvider(messageConverterProvider, 1, deviceClientProvider.Object, Option.None<UpstreamProtocol>());
             IConnectionManager connectionManager = new ConnectionManager(cloudConnectionProvider);
 
-            IClientCredentials clientCredentials1 = GetClientCredentials();
+            IClientCredentials clientCredentials1 = GetClientCredentials(TimeSpan.FromSeconds(10));
             Try<ICloudProxy> cloudProxyTry1 = await connectionManager.CreateCloudConnectionAsync(clientCredentials1);
             Assert.True(cloudProxyTry1.Success);
 
@@ -352,36 +352,26 @@ IClient GetMockedDeviceClient()
             Task<string> tokenGetter = deviceTokenRefresher.GetTokenAsync(hostname);
             Assert.False(tokenGetter.IsCompleted);
 
-            IClientCredentials clientCredentials2 = GetClientCredentials();
+            IClientCredentials clientCredentials2 = GetClientCredentials(TimeSpan.FromMinutes(2));
             Try<ICloudProxy> cloudProxyTry2 = await connectionManager.CreateCloudConnectionAsync(clientCredentials2);
             Assert.True(cloudProxyTry2.Success);
 
             IDeviceProxy deviceProxy2 = GetMockDeviceProxy();
             await connectionManager.AddDeviceConnection(clientCredentials2.Identity, deviceProxy2);
 
             await Task.Delay(TimeSpan.FromSeconds(3));
-            Assert.True(tokenGetter.IsCompleted);
-            Assert.Equal(tokenGetter.Result, (clientCredentials2 as ITokenCredentials)?.Token);
-
-            await Task.Delay(TimeSpan.FromSeconds(10));
-            Assert.NotNull(authenticationMethod);
-            deviceTokenRefresher = authenticationMethod as DeviceAuthenticationWithTokenRefresh;
-            Assert.NotNull(deviceTokenRefresher);
-            tokenGetter = deviceTokenRefresher.GetTokenAsync(hostname);
             Assert.False(tokenGetter.IsCompleted);
 
-            IClientCredentials clientCredentials3 = GetClientCredentials();
+            IClientCredentials clientCredentials3 = GetClientCredentials(TimeSpan.FromMinutes(10));
             Try<ICloudProxy> cloudProxyTry3 = await connectionManager.CreateCloudConnectionAsync(clientCredentials3);
             Assert.True(cloudProxyTry3.Success);
 
             IDeviceProxy deviceProxy3 = GetMockDeviceProxy();
             await connectionManager.AddDeviceConnection(clientCredentials3.Identity, deviceProxy3);
 
-            await Task.Delay(TimeSpan.FromSeconds(3));
+            await Task.Delay(TimeSpan.FromSeconds(23));
             Assert.True(tokenGetter.IsCompleted);
             Assert.Equal(tokenGetter.Result, (clientCredentials3 as ITokenCredentials)?.Token);
-
-            Mock.VerifyAll(Mock.Get(deviceProxy1), Mock.Get(deviceProxy2));
         }
 
         static async Task GetCloudConnectionTest(Func<IClientCredentials> credentialsGenerator, IClientProvider clientProvider)