Skip to content

Commit e90adbf

Browse files
msJinLeimsJinLei
committed
Add test cases to collect authentication telemetry
1 parent c293e6b commit e90adbf

File tree

5 files changed

+337
-180
lines changed

5 files changed

+337
-180
lines changed

‎src/Accounts/Authentication.Test/AuthenticationFactoryConcurrencyTest.cs

+246
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,246 @@
1+
// ----------------------------------------------------------------------------------
2+
//
3+
// Copyright Microsoft Corporation
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
// ----------------------------------------------------------------------------------
14+
using Azure.Core;
15+
using Azure.Identity;
16+
17+
using Microsoft.Azure.Commands.Common.Authentication;
18+
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
19+
using Microsoft.Azure.Commands.Common.Authentication.Factories;
20+
using Microsoft.Azure.Commands.Common.Authentication.Test.Mocks;
21+
using Microsoft.Azure.Commands.ResourceManager.Common;
22+
using Microsoft.Azure.PowerShell.Authenticators.Factories;
23+
using Microsoft.WindowsAzure.Commands.Common;
24+
25+
using Moq;
26+
27+
using System;
28+
using System.Collections.Generic;
29+
using System.Linq;
30+
using System.Security;
31+
using System.Threading;
32+
using System.Threading.Tasks;
33+
34+
using Xunit;
35+
using Xunit.Abstractions;
36+
37+
namespace Common.Authentication.Test
38+
{
39+
public class AuthenticationFactoryConcurrencyTest : MockAuthenticationTestBase
40+
{
41+
public string tenant = Guid.NewGuid().ToString();
42+
public string armToken = Guid.NewGuid().ToString();
43+
44+
public AuthenticationFactoryConcurrencyTest(ITestOutputHelper output) : base(output)
45+
{
46+
47+
}
48+
49+
[Fact]
50+
public void AuthenticationFactoryShouldHaveTelemetry()
51+
{
52+
//Set the inputs
53+
var account = new AzureAccount
54+
{
55+
Id = Guid.NewGuid().ToString(),
56+
Type = AzureAccount.AccountType.ServicePrincipal
57+
};
58+
account.SetTenants(tenant);
59+
var environment = AzureEnvironment.PublicEnvironments.Values.First();
60+
var password = "password".ConvertToSecureString();
61+
var promptBehavior = "auto";
62+
Action<string> promptAction = s => { };
63+
var resourceId = "resourceId";
64+
var cmdletContext = new AzureCmdletContext(Guid.NewGuid().ToString());
65+
66+
var keyStore = new Dictionary<IKeyStoreKey, SecureString>();
67+
68+
// Arrange
69+
var factory = new AuthenticationFactory();
70+
71+
// Mock AzureSession setup
72+
try
73+
{
74+
originalInstance = AzureSession.Instance;
75+
}
76+
catch (InvalidOperationException)
77+
{
78+
originalInstance = null;
79+
}
80+
try
81+
{
82+
InitializeSession(armToken);
83+
var mockTokenCredential = new Mock<ClientSecretCredential>();
84+
mockTokenCredential.Setup(f => f.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>())).ReturnsAsync(new AccessToken(armToken, DateTimeOffset.Now));
85+
mockTokenCredential.Setup(f => f.GetToken(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>())).Returns(new AccessToken(armToken, DateTimeOffset.Now));
86+
mockAzureCredentialFactory.Setup(f => f.CreateClientSecretCredential(
87+
It.IsAny<string>(),
88+
It.IsAny<string>(),
89+
It.IsAny<SecureString>(),
90+
It.IsAny<ClientSecretCredentialOptions>())
91+
).Returns(mockTokenCredential.Object);
92+
AzureSession.Instance.RegisterComponent(nameof(AzureCredentialFactory), () => mockAzureCredentialFactory.Object);
93+
94+
mockKeyStore.Setup(f => f.SaveSecureString(It.IsAny<IKeyStoreKey>(), It.IsAny<SecureString>())).Callback((IKeyStoreKey k, SecureString v) => keyStore[k] = v);
95+
mockKeyStore.Setup(f => f.GetSecureString(It.IsAny<IKeyStoreKey>())).Returns(password);
96+
AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => mockKeyStore.Object);
97+
98+
// Act
99+
var optionalParameters = new Dictionary<string, object>()
100+
{
101+
{AuthenticationFactory.TokenCacheParameterName, mockTokenCache.Object},
102+
{AuthenticationFactory.ResourceIdParameterName, resourceId },
103+
{AuthenticationFactory.CmdletContextParameterName, cmdletContext }
104+
};
105+
var result = factory.Authenticate(account, environment, tenant, password, promptBehavior, promptAction, optionalParameters);
106+
107+
// Assert
108+
Assert.NotNull(result);
109+
Assert.Equal(armToken, result.AccessToken);
110+
Assert.Equal(account.Id, result.UserId);
111+
112+
if (AzureSession.Instance.TryGetComponent(AuthenticationTelemetry.Name, out AuthenticationTelemetry authenticationTelemetry))
113+
{
114+
var telemetry = authenticationTelemetry.GetTelemetryRecord(cmdletContext);
115+
Assert.NotNull(telemetry.Primary);
116+
Assert.Equal(mockTokenCredential.Object.GetType().Name, telemetry.Primary.TokenCredentialName);
117+
Assert.True(telemetry.Primary.AuthenticationSuccess);
118+
Assert.Empty(telemetry.Secondary);
119+
}
120+
else
121+
{
122+
Assert.Fail("Authentication telemetry component not found in AzureSession");
123+
}
124+
}
125+
finally
126+
{
127+
// Restore original instance
128+
AzureSession.Initialize(() => originalInstance, true);
129+
}
130+
}
131+
132+
[Fact]
133+
public async Task AuthenticationFactoryShouldHandleConcurrentAuthentication()
134+
{
135+
// Set up common test data
136+
var environment = AzureEnvironment.PublicEnvironments.Values.First();
137+
var promptBehavior = "auto";
138+
Action<string> promptAction = s => { };
139+
var resourceId = "resourceId";
140+
141+
var keyStore = new Dictionary<IKeyStoreKey, SecureString>();
142+
var taskCount = 5;
143+
144+
// Arrange
145+
var factory = new AuthenticationFactory();
146+
147+
// Mock AzureSession setup
148+
try
149+
{
150+
originalInstance = AzureSession.Instance;
151+
}
152+
catch (InvalidOperationException)
153+
{
154+
originalInstance = null;
155+
}
156+
157+
try
158+
{
159+
InitializeSession(armToken);
160+
var mockTokenCredential = new Mock<SharedTokenCacheCredential>();
161+
mockTokenCredential.Setup(f => f.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>()))
162+
.ReturnsAsync(new AccessToken(armToken, DateTimeOffset.Now));
163+
mockAzureCredentialFactory.Setup(f => f.CreateSharedTokenCacheCredentials(It.IsAny<SharedTokenCacheCredentialOptions>())).Returns(mockTokenCredential.Object);
164+
165+
//var mockTokenManagedIdentityCredential = new Mock<ManagedIdentityCredential>();
166+
//mockTokenManagedIdentityCredential.Setup(f => f.GetTokenAsync(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>()))
167+
// .ReturnsAsync(new AccessToken(armToken, DateTimeOffset.Now));
168+
//mockTokenManagedIdentityCredential.Setup(f => f.GetToken(It.IsAny<TokenRequestContext>(), It.IsAny<CancellationToken>()))
169+
// .Returns(new AccessToken(armToken, DateTimeOffset.Now));
170+
//mockAzureCredentialFactory.Setup(f => f.CreateManagedIdentityCredential(It.IsAny<string>())).Returns(mockTokenManagedIdentityCredential.Object);
171+
172+
AzureSession.Instance.RegisterComponent(nameof(AzureCredentialFactory), () => mockAzureCredentialFactory.Object);
173+
174+
// Act - Create multiple tasks to authenticate concurrently
175+
var tasks = new Task<IAccessToken>[taskCount];
176+
var contexts = new AzureCmdletContext[taskCount];
177+
178+
for (int i = 0; i < taskCount; i++)
179+
{
180+
// Create unique accounts and context for each task
181+
var account = new AzureAccount
182+
{
183+
Id = Guid.NewGuid().ToString(),
184+
Type = AzureAccount.AccountType.User
185+
};
186+
account.SetTenants(tenant);
187+
contexts[i] = new AzureCmdletContext(Guid.NewGuid().ToString());
188+
189+
var optionalParameters = new Dictionary<string, object>()
190+
{
191+
{AuthenticationFactory.TokenCacheParameterName, mockTokenCache.Object},
192+
{AuthenticationFactory.ResourceIdParameterName, resourceId},
193+
{AuthenticationFactory.CmdletContextParameterName, contexts[i]}
194+
};
195+
196+
// Capture variables for the lambda to avoid closure issues
197+
var currentAccount = account;
198+
var currentContext = contexts[i];
199+
var taskIndex = i;
200+
201+
tasks[i] = Task.Run(() =>
202+
{
203+
_output.WriteLine($"Starting authentication task {taskIndex}");
204+
var result = factory.Authenticate(currentAccount, environment, tenant,
205+
null, promptBehavior, promptAction, optionalParameters);
206+
_output.WriteLine($"Completed authentication task {taskIndex}");
207+
return result;
208+
});
209+
}
210+
211+
// Wait for all tasks to complete
212+
await Task.WhenAll(tasks);
213+
214+
// Collect all results
215+
for (int i = 0; i < taskCount; i++)
216+
{
217+
var result = await tasks[i];
218+
Assert.NotNull(result);
219+
Assert.Equal(armToken, result.AccessToken);
220+
}
221+
222+
// Verify telemetry was recorded for each authentication
223+
if (AzureSession.Instance.TryGetComponent(AuthenticationTelemetry.Name, out AuthenticationTelemetry authenticationTelemetry))
224+
{
225+
for (int i = 0; i < taskCount; i++)
226+
{
227+
var telemetry = authenticationTelemetry.GetTelemetryRecord(contexts[i]);
228+
Assert.NotNull(telemetry.Primary);
229+
Assert.Equal(mockTokenCredential.Object.GetType().Name, telemetry.Primary.TokenCredentialName);
230+
Assert.True(telemetry.Primary.AuthenticationSuccess);
231+
Assert.Empty(telemetry.Secondary);
232+
}
233+
}
234+
else
235+
{
236+
Assert.Fail("Authentication telemetry component not found in AzureSession");
237+
}
238+
}
239+
finally
240+
{
241+
// Restore original instance
242+
AzureSession.Initialize(() => originalInstance, true);
243+
}
244+
}
245+
}
246+
}

‎src/Accounts/Authentication.Test/AuthenticationFactoryTests.cs

+11-15
Original file line numberDiff line numberDiff line change
@@ -12,32 +12,28 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15-
using System;
16-
using System.Collections.Generic;
17-
using System.Linq;
15+
using Azure.Core;
1816

1917
using Microsoft.Azure.Commands.Common.Authentication;
2018
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
2119
using Microsoft.Azure.Commands.Common.Authentication.Factories;
2220
using Microsoft.Azure.Commands.Common.Authentication.Test;
21+
using Microsoft.Azure.Commands.TestFx.Mocks;
22+
using Microsoft.Azure.PowerShell.Authentication.Test.Mocks;
2323
using Microsoft.Azure.PowerShell.Authenticators;
24+
using Microsoft.Azure.PowerShell.Authenticators.Factories;
2425
using Microsoft.WindowsAzure.Commands.ScenarioTest;
25-
using Microsoft.WindowsAzure.Commands.Test.Utilities.Common;
2626
using Microsoft.WindowsAzure.Commands.Utilities.Common;
2727

28-
using Xunit;
29-
using Xunit.Abstractions;
30-
using System.Text.RegularExpressions;
28+
using System;
29+
using System.Collections.Generic;
30+
using System.Linq;
3131
using System.Net.Http;
32+
using System.Text.RegularExpressions;
3233
using System.Threading;
33-
using Microsoft.Azure.PowerShell.Authenticators.Factories;
34-
using Microsoft.WindowsAzure.Commands.Common.Test.Mocks;
35-
using Microsoft.Azure.PowerShell.Authentication.Test.Mocks;
36-
using Azure.Identity;
37-
using Moq;
38-
using System.ServiceModel.Channels;
39-
using Azure.Core;
40-
using Microsoft.Azure.Commands.TestFx.Mocks;
34+
35+
using Xunit;
36+
using Xunit.Abstractions;
4137

4238
namespace Common.Authentication.Test
4339
{

‎src/Accounts/Authentication.Test/Mocks/MockAuthenticationTestBase.cs

+76
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
// ----------------------------------------------------------------------------------
2+
//
3+
// Copyright Microsoft Corporation
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
// ----------------------------------------------------------------------------------
14+
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
15+
using Microsoft.Azure.Commands.Common.Authentication.Models;
16+
using Microsoft.Azure.Commands.ResourceManager.Common;
17+
using Microsoft.Azure.PowerShell.Authenticators;
18+
using Microsoft.Azure.PowerShell.Authenticators.Factories;
19+
20+
using Moq;
21+
22+
using System;
23+
using System.IO;
24+
25+
using Xunit.Abstractions;
26+
27+
using static Microsoft.Azure.Commands.Common.Authentication.AzureSessionInitializer;
28+
29+
namespace Microsoft.Azure.Commands.Common.Authentication.Test.Mocks
30+
{
31+
public abstract class MockAuthenticationTestBase
32+
{
33+
protected Mock<AzKeyStore> mockKeyStore;
34+
protected Mock<IAzureTokenCache> mockTokenCache;
35+
protected Mock<AzureCredentialFactory> mockAzureCredentialFactory;
36+
protected IAzureSession originalInstance;
37+
38+
public ITestOutputHelper _output;
39+
public MockAuthenticationTestBase(ITestOutputHelper output)
40+
{
41+
_output = output;
42+
43+
mockTokenCache = new Mock<IAzureTokenCache>();
44+
mockKeyStore = new Mock<AzKeyStore>(string.Empty, string.Empty, true, null);
45+
mockAzureCredentialFactory = new Mock<AzureCredentialFactory>();
46+
}
47+
48+
static IAzureSession CreateInstance()
49+
{
50+
var profilePath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), ".Azure");
51+
var dataStore = new MemoryDataStore();
52+
var session = new AdalSession
53+
{
54+
DataStore = dataStore,
55+
ProfileDirectory = profilePath,
56+
ProfileFile = "AzureProfile.json",
57+
TokenCacheDirectory = MsalCacheHelperProvider.MsalTokenCachePath,
58+
TokenCacheFile = MsalCacheHelperProvider.GetTokenCacheName(MsalCacheHelperProvider.LegacyTokenCacheName, caeEnabled: true)
59+
};
60+
61+
session.TokenCache = session.TokenCache ?? new AzureTokenCache();
62+
return session;
63+
}
64+
65+
public void InitializeSession(string dummyToken)
66+
{
67+
var mockAzureSession = CreateInstance();
68+
AzureSession.Initialize(() => mockAzureSession, true);
69+
70+
AzureSession.Instance.RegisterComponent(AuthenticatorBuilder.AuthenticatorBuilderKey, () => (IAuthenticatorBuilder)new DefaultAuthenticatorBuilder());
71+
AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => (PowerShellTokenCacheProvider)new InMemoryTokenCacheProvider());
72+
AzureSession.Instance.RegisterComponent(nameof(MsalAccessTokenAcquirerFactory), () => new MsalAccessTokenAcquirerFactory());
73+
AzureSession.Instance.RegisterComponent<AuthenticationTelemetry>(AuthenticationTelemetry.Name, () => new AuthenticationTelemetry());
74+
}
75+
}
76+
}

0 commit comments

Comments
 (0)