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
+ }
0 commit comments