From 38a3dca1e089970e32706758a1c50b2d6126fd5a Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Mar 2022 14:15:32 +0000 Subject: [PATCH] Reacted to code review feedback --- .../CosmosItemSessionTokenTests.cs | 27 +++++-------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemSessionTokenTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemSessionTokenTests.cs index 8d96d2b8ec..9b77802c01 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemSessionTokenTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemSessionTokenTests.cs @@ -5,9 +5,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests { using System; - using System.Collections.Generic; - using System.Globalization; - using System.Linq; using System.Threading.Tasks; using Microsoft.Azure.Documents; using Microsoft.Azure.Cosmos; @@ -16,6 +13,8 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests using Microsoft.Azure.Cosmos.Diagnostics; using System.IO; using System.Net; + using System.Threading; + using Microsoft.Azure.Cosmos.Routing; [TestClass] public class CosmosItemSessionTokenTests : BaseCosmosClientHelper @@ -235,26 +234,12 @@ public async Task NoSessionTokenCaptureForThrottledUpsertRequestsTest() Assert.IsNull(lsnAfterThrottledRequest); } - private async Task GetPKRangeIdForPartitionKey( - PartitionKey pkValue) + private async Task GetPKRangeIdForPartitionKey(PartitionKey pkValue) { - DocumentFeedResponse pkRanges = await this.cosmosClient.DocumentClient.ReadPartitionKeyRangeFeedAsync( - UriFactory.CreateDocumentCollectionUri(this.Container.Database.Id, this.Container.Id)); - List maxExclusiveBoundaries = pkRanges.Select(pkRange => pkRange.MaxExclusive).ToList(); + CollectionRoutingMap collectionRoutingMap = await ((ContainerInternal)this.Container).GetRoutingMapAsync(CancellationToken.None); + string effectivePK = pkValue.InternalKey.GetEffectivePartitionKeyString(this.containerSettings.PartitionKey); - string effectivePK1 = pkValue.InternalKey.GetEffectivePartitionKeyString(this.containerSettings.PartitionKey); - int pkIndex = 0; - while (pkIndex < maxExclusiveBoundaries.Count && string.Compare(effectivePK1, maxExclusiveBoundaries[pkIndex]) >= 0) - { - ++pkIndex; - } - - if (pkIndex == maxExclusiveBoundaries.Count) - { - throw new Exception("Failed to find the range"); - } - - return pkIndex.ToString(CultureInfo.InvariantCulture); + return collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePK).Id; } private async Task> GetLSNFromSessionContainer(PartitionKey pkValue)