@@ -60,7 +60,7 @@ public class ManagedIdentityTests
6060 //non existent Resource ID of the User Assigned Identity
6161 private const string Non_Existent_UamiResourceId = "/subscriptions/userAssignedIdentities/NO_ID" ;
6262
63- // [DataTestMethod]
63+ [ DataTestMethod ]
6464 [ DataRow ( MsiAzureResource . WebApp , "" , DisplayName = "System_Identity_Web_App" ) ]
6565 //[DataRow(MsiAzureResource.Function, "", DisplayName = "System_Identity_Function_App")]
6666 //[DataRow(MsiAzureResource.VM, "", DisplayName = "System_Identity_Virtual_Machine")]
@@ -128,7 +128,69 @@ public async Task AcquireMSITokenAsync(MsiAzureResource azureResource, string us
128128 }
129129 }
130130
131- //[TestMethod]
131+ [ DataTestMethod ]
132+ [ DataRow ( MsiAzureResource . WebApp , "" , DisplayName = "System_Identity_Web_App" ) ]
133+ [ DataRow ( MsiAzureResource . WebApp , UserAssignedClientID , UserAssignedIdentityId . ClientId , DisplayName = "ClientId_Web_App" ) ]
134+ [ DataRow ( MsiAzureResource . WebApp , UamiResourceId , UserAssignedIdentityId . ResourceId , DisplayName = "ResourceID_Web_App" ) ]
135+ [ DataRow ( MsiAzureResource . WebApp , UserAssignedObjectID , UserAssignedIdentityId . ObjectId , DisplayName = "ObjectID_Web_App" ) ]
136+ public async Task AcquireMSITokenWithClaimsAsync (
137+ MsiAzureResource azureResource ,
138+ string userIdentity ,
139+ UserAssignedIdentityId userAssignedIdentityId = UserAssignedIdentityId . None )
140+ {
141+ using ( new EnvVariableContext ( ) )
142+ {
143+ // ---------- Arrange ----------
144+ var envVariables = await GetEnvironmentVariablesAsync ( azureResource ) . ConfigureAwait ( false ) ;
145+ SetEnvironmentVariables ( envVariables ) ;
146+
147+ string uri = s_baseURL + $ "MSIToken?azureresource={ azureResource } &uri=";
148+
149+ IManagedIdentityApplication mia =
150+ CreateMIAWithProxy ( uri , userIdentity , userAssignedIdentityId ) ;
151+
152+ // ---------- Act & Assert 1 ----------
153+ AuthenticationResult result1 = await mia
154+ . AcquireTokenForManagedIdentity ( s_msi_scopes )
155+ . ExecuteAsync ( )
156+ . ConfigureAwait ( false ) ;
157+
158+ Assert . AreEqual ( "Bearer" , result1 . TokenType ) ;
159+ Assert . AreEqual ( TokenSource . IdentityProvider ,
160+ result1 . AuthenticationResultMetadata . TokenSource ) ;
161+ CoreAssert . IsWithinRange (
162+ DateTimeOffset . UtcNow ,
163+ result1 . ExpiresOn ,
164+ TimeSpan . FromHours ( 24 ) ) ;
165+
166+ // ---------- Act & Assert 2 (cache hit) ----------
167+ AuthenticationResult result2 = await mia
168+ . AcquireTokenForManagedIdentity ( s_msi_scopes )
169+ . ExecuteAsync ( )
170+ . ConfigureAwait ( false ) ;
171+
172+ Assert . IsTrue ( result2 . Scopes . All ( s_msi_scopes . Contains ) ) ;
173+ Assert . AreEqual ( TokenSource . Cache ,
174+ result2 . AuthenticationResultMetadata . TokenSource ) ;
175+ Assert . AreEqual ( result1 . AccessToken , result2 . AccessToken , // sanity
176+ "Second call should come from cache" ) ;
177+
178+ // ---------- Act & Assert 3 (claims → bypass_cache) ----------
179+ const string claimsJson = TestConstants . Claims ;
180+
181+ AuthenticationResult result3 = await mia
182+ . AcquireTokenForManagedIdentity ( s_msi_scopes )
183+ . WithClaims ( claimsJson )
184+ . ExecuteAsync ( )
185+ . ConfigureAwait ( false ) ;
186+
187+ // Token source should now be IdentityProvider again
188+ Assert . AreEqual ( TokenSource . IdentityProvider ,
189+ result3 . AuthenticationResultMetadata . TokenSource ) ;
190+ }
191+ }
192+
193+ [ TestMethod ]
132194 public async Task AcquireMsiToken_ForTokenExchangeResource_Successfully ( )
133195 {
134196 string resource = "api://AzureAdTokenExchange" ;
@@ -183,7 +245,7 @@ public async Task AcquireMsiToken_ForTokenExchangeResource_Successfully()
183245 }
184246 }
185247
186- // [TestMethod]
248+ [ TestMethod ]
187249 public async Task AcquireMsiToken_ExchangeForEstsToken_Successfully ( )
188250 {
189251 const string resource = "api://AzureAdTokenExchange" ;
@@ -449,7 +511,7 @@ private IManagedIdentityApplication CreateMIAWithProxy(string url, string userAs
449511 // Disabling shared cache options to avoid cross test pollution.
450512 builder . Config . AccessorOptions = null ;
451513
452- IManagedIdentityApplication mia = builder
514+ IManagedIdentityApplication mia = builder . WithClientCapabilities ( new [ ] { "cp1" } )
453515 . WithHttpManager ( proxyHttpManager ) . Build ( ) ;
454516
455517 return mia ;
0 commit comments