changeToUseChangeFeedToFetchPkRanges

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/InCompleteRoutingMapRetryPolicyTest.java

@@ -0,0 +1,55 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.implementation.caches.InCompleteRoutingMapRetryPolicy;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import reactor.test.StepVerifier;
+
+public class InCompleteRoutingMapRetryPolicyTest {
+    private InCompleteRoutingMapRetryPolicy retryPolicy;
+
+    @BeforeMethod(groups = {"unit"})
+    public void beforeMethod() {
+        retryPolicy = new InCompleteRoutingMapRetryPolicy();
+    }
+
+    @Test(groups = {"unit"})
+    public void shouldRetry_WithInCompleteRoutingMapException() {
+        InCompleteRoutingMapException exception = new InCompleteRoutingMapException("test");
+
+        // First attempt - should retry
+        StepVerifier.create(retryPolicy.shouldRetry(exception))
+            .expectNext(ShouldRetryResult.RETRY_NOW)
+            .verifyComplete();
+
+        // Second attempt - should not retry
+        StepVerifier.create(retryPolicy.shouldRetry(exception))
+            .expectNext(ShouldRetryResult.NO_RETRY)
+            .verifyComplete();
+    }
+
+    @Test(groups = {"unit"})
+    public void shouldRetry_WithOtherException_ShouldNotRetry() {
+        Exception otherException = new RuntimeException("Some other error");
+
+        StepVerifier.create(retryPolicy.shouldRetry(otherException))
+            .expectNext(ShouldRetryResult.NO_RETRY)
+            .verifyComplete();
+    }
+
+    @Test(groups = {"unit"})
+    public void shouldRetry_WithNullException_ShouldNotRetry() {
+        StepVerifier.create(retryPolicy.shouldRetry(null))
+            .expectNext(ShouldRetryResult.NO_RETRY)
+            .verifyComplete();
+    }
+
+    @Test(groups = {"unit"})
+    public void getRetryContext_ReturnsNull() {
+        RetryContext context = retryPolicy.getRetryContext();
+        assert context == null : "RetryContext should be null";
+    }
+}

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/RxDocumentClientImplTest.java

@@ -394,7 +394,15 @@ public String getCollectionUniqueId() {
             }
 
             @Override
-            public CollectionRoutingMap tryCombine(List<ImmutablePair<PartitionKeyRange, IServerIdentity>> ranges) {
+            public CollectionRoutingMap tryCombine(
+                List<ImmutablePair<PartitionKeyRange, IServerIdentity>> ranges,
+                String changeFeedIfNoneMatch,
+                String collectionRid) {
+                return null;
+            }
+
+            @Override
+            public String getChangeFeedNextIfNoneMatch() {
                 return null;
             }
         };

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/caches/RxPartitionKeyRangeCacheTest.java

@@ -0,0 +1,248 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation.caches;
+
+import com.azure.cosmos.implementation.DocumentCollection;
+import com.azure.cosmos.implementation.InCompleteRoutingMapException;
+import com.azure.cosmos.implementation.PartitionKeyRange;
+import com.azure.cosmos.implementation.RxDocumentClientImpl;
+import com.azure.cosmos.implementation.Utils;
+import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
+import com.azure.cosmos.implementation.routing.CollectionRoutingMap;
+import com.azure.cosmos.implementation.routing.InMemoryCollectionRoutingMap;
+import com.azure.cosmos.models.CosmosQueryRequestOptions;
+import com.azure.cosmos.models.FeedResponse;
+import org.mockito.Mockito;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+
+import static org.assertj.core.api.Fail.fail;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+public class RxPartitionKeyRangeCacheTest {
+    private RxDocumentClientImpl client;
+    private RxCollectionCache collectionCache;
+    private RxPartitionKeyRangeCache cache;
+
+    @BeforeMethod(groups = "unit")
+    public void before_test() {
+        client = Mockito.mock(RxDocumentClientImpl.class);
+        collectionCache = Mockito.mock(RxCollectionCache.class);
+        cache = new RxPartitionKeyRangeCache(client, collectionCache);
+    }
+
+    @Test(groups = "unit")
+    public void getRoutingMapUsesChangeFeedNextIfNoneMatchWhenNotEmpty() {
+        String collectionRid = "collection1";
+        String changeFeedToken = "token1";
+
+        PartitionKeyRange range1 = new PartitionKeyRange();
+        range1.setId("0");
+        range1.setMinInclusive(PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY);
+        range1.setMaxExclusive(PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY);
+
+        CollectionRoutingMap previousRoutingMap = InMemoryCollectionRoutingMap
+            .tryCreateCompleteRoutingMap(Arrays.asList(ImmutablePair.of(range1, null)), collectionRid, changeFeedToken);
+
+        DocumentCollection collection = new DocumentCollection();
+        collection.setResourceId(collectionRid);
+        collection.setSelfLink("dbs/db1/colls/coll1");
+
+        FeedResponse<PartitionKeyRange> response = Mockito.mock(FeedResponse.class);
+        when(response.getResults()).thenReturn(Arrays.asList(range1));
+        when(response.getContinuationToken()).thenReturn("newToken");
+
+        when(collectionCache.resolveCollectionAsync(any(), any()))
+            .thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
+
+        when(client.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
+            .thenReturn(Flux.just(response));
+
+        StepVerifier.create(cache.tryLookupAsync(null, collectionRid, previousRoutingMap, new HashMap<>()))
+            .expectNextMatches(routingMapHolder -> 
+                routingMapHolder != null && 
+                routingMapHolder.v != null &&
+                changeFeedToken.equals(previousRoutingMap.getChangeFeedNextIfNoneMatch()))
+            .verifyComplete();
+    }
+
+    @Test(groups = "unit")
+    public void getRoutingMapWithEmptyChangeFeedNextIfNoneMatch() {
+        String collectionRid = "collection1";
+
+        PartitionKeyRange range1 = new PartitionKeyRange();
+        range1.setId("0");
+        range1.setMinInclusive(PartitionKeyRange.MINIMUM_INCLUSIVE_EFFECTIVE_PARTITION_KEY);
+        range1.setMaxExclusive(PartitionKeyRange.MAXIMUM_EXCLUSIVE_EFFECTIVE_PARTITION_KEY);
+
+        CollectionRoutingMap previousRoutingMap = InMemoryCollectionRoutingMap
+            .tryCreateCompleteRoutingMap(
+                Arrays.asList(ImmutablePair.of(range1, null)),
+                collectionRid,
+                null);
+
+        DocumentCollection collection = new DocumentCollection();
+        collection.setResourceId(collectionRid);
+        collection.setSelfLink("dbs/db1/colls/coll1");
+
+        FeedResponse<PartitionKeyRange> response = Mockito.mock(FeedResponse.class);
+        when(response.getResults()).thenReturn(Arrays.asList(range1));
+        when(response.getContinuationToken()).thenReturn("newToken");
+
+        when(collectionCache.resolveCollectionAsync(any(), any()))
+            .thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
+
+        when(client.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
+            .thenReturn(Flux.just(response));
+
+        StepVerifier.create(cache.tryLookupAsync(null, collectionRid, previousRoutingMap, new HashMap<>()))
+            .expectNextMatches(routingMapHolder ->
+                routingMapHolder != null &&
+                routingMapHolder.v != null &&
+                previousRoutingMap.getChangeFeedNextIfNoneMatch() == null)
+            .verifyComplete();
+    }
+
+    @Test(groups = "unit")
+    public void validateIsCompleteSetOfRanges_ValidRanges() {
+        PartitionKeyRange range1 = new PartitionKeyRange();
+        range1.setId("0");
+        range1.setMinInclusive("");
+        range1.setMaxExclusive("0000000030");
+
+        PartitionKeyRange range2 = new PartitionKeyRange();
+        range2.setId("1");
+        range2.setMinInclusive("0000000030");
+        range2.setMaxExclusive("0000000070");
+
+        PartitionKeyRange range3 = new PartitionKeyRange();
+        range3.setId("2");
+        range3.setMinInclusive("0000000070");
+        range3.setMaxExclusive("FF");
+
+        CollectionRoutingMap routingMap = InMemoryCollectionRoutingMap
+            .tryCreateCompleteRoutingMap(
+                Arrays.asList(
+                    new ImmutablePair<>(range1, null),
+                    new ImmutablePair<>(range2, null),
+                    new ImmutablePair<>(range3, null)
+                ),
+                "dummyCollectionId",
+                null);
+
+        // Verify that the routing map was created successfully
+        // A non-null result indicates the ranges form a complete set
+        assertNotNull(routingMap, "Routing map should be created for valid complete set of ranges");
+    }
+
+    @Test(groups = "unit")
+    public void validateIsCompleteSetOfRanges_WithOverlap() {
+        PartitionKeyRange range1 = new PartitionKeyRange();
+        range1.setId("0");
+        range1.setMinInclusive("");
+        range1.setMaxExclusive("0000000050");
+
+        PartitionKeyRange range2 = new PartitionKeyRange();
+        range2.setId("1");
+        range2.setMinInclusive("0000000030"); // Overlaps with range1
+        range2.setMaxExclusive("FF");
+
+        try {
+            InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(
+                Arrays.asList(
+                    new ImmutablePair<>(range1, null),
+                    new ImmutablePair<>(range2, null)
+                ),
+                "dummyCollectionId",
+                null);
+            fail("Expected InCompleteRoutingMapException for overlapping ranges");
+        } catch (InCompleteRoutingMapException e) {
+            assertEquals(
+                e.getMessage(),
+                "Ranges overlap for collectionRid dummyCollectionId, previous range [{\"min\":\"\",\"max\":\"0000000050\"}], current range [{\"min\":\"0000000030\",\"max\":\"FF\"}]",
+                "Unexpected error message");
+        }
+    }
+
+    @Test(groups = "unit")
+    public void validateIsCompleteSetOfRanges_EmptyList() {
+        try {
+            InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(
+                Collections.emptyList(),
+                "dummyCollectionId",
+                null);
+            fail("Expected InCompleteRoutingMapException for empty ranges list");
+        } catch (InCompleteRoutingMapException e) {
+            assertEquals(e.getMessage(), "Empty ranges for collectionRid dummyCollectionId", "Unexpected error message");
+        }
+    }
+
+    @Test(groups = "unit")
+    public void validateIsCompleteSetOfRanges_WithGap() {
+        PartitionKeyRange range1 = new PartitionKeyRange();
+        range1.setId("0");
+        range1.setMinInclusive("");
+        range1.setMaxExclusive("0000000020");
+
+        PartitionKeyRange range2 = new PartitionKeyRange();
+        range2.setId("1");
+        range2.setMinInclusive("0000000030"); // Overlaps with range1
+        range2.setMaxExclusive("FF");
+
+        try {
+            InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(
+                Arrays.asList(
+                    new ImmutablePair<>(range1, null),
+                    new ImmutablePair<>(range2, null)
+                ),
+                "dummyCollectionId",
+                null);
+            fail("Expected InCompleteRoutingMapException for overlapping ranges");
+        } catch (InCompleteRoutingMapException e) {
+            assertEquals(
+                e.getMessage(),
+                "Ranges incomplete for collectionRid dummyCollectionId, previous range [{\"min\":\"\",\"max\":\"0000000020\"}], current range [{\"min\":\"0000000030\",\"max\":\"FF\"}]",
+                "Unexpected error message");
+        }
+    }
+
+    @Test(groups = "unit")
+    public void tryLookupAsync_RetriesOnceAndConvertsToNotFoundException() {
+        String collectionRid = "collection1";
+
+        DocumentCollection collection = new DocumentCollection();
+        collection.setResourceId(collectionRid);
+        collection.setSelfLink("dbs/db1/colls/coll1");
+
+        // First attempt - returns incomplete routing map
+        FeedResponse<PartitionKeyRange> response1 = Mockito.mock(FeedResponse.class);
+        when(response1.getResults()).thenReturn(Collections.emptyList());
+
+        // Second attempt - still returns incomplete routing map
+        FeedResponse<PartitionKeyRange> response2 = Mockito.mock(FeedResponse.class);
+        when(response2.getResults()).thenReturn(Collections.emptyList());
+
+        when(collectionCache.resolveCollectionAsync(any(), any()))
+            .thenReturn(Mono.just(new Utils.ValueHolder<>(collection)));
+
+        when(client.readPartitionKeyRanges(eq(collection.getSelfLink()), any(CosmosQueryRequestOptions.class)))
+            .thenReturn(Flux.just(response1))
+            .thenReturn(Flux.just(response2));
+
+        StepVerifier.create(cache.tryLookupAsync(null, collection.getResourceId(), null, new HashMap<>()))
+            .expectNextMatches(s -> s.v == null)
+            .verifyComplete();
+    }
+}

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/directconnectivity/AddressResolverTest.java

@@ -108,7 +108,8 @@ public void before_AddressResolverTest() throws Exception {
         this.routingMapCollection1BeforeSplit =
             InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(
                 rangesBeforeSplit1,
-                collection1.getResourceId());
+                collection1.getResourceId(),
+                "1");
 
         List<ImmutablePair<PartitionKeyRange, IServerIdentity>> rangesAfterSplit1 =
             new ArrayList<>();
@@ -127,7 +128,7 @@ public void before_AddressResolverTest() throws Exception {
 
         addPartitionKeyRangeFunc.apply(rangesAfterSplit1);
 
-        this.routingMapCollection1AfterSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesAfterSplit1, collection1.getResourceId());
+        this.routingMapCollection1AfterSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesAfterSplit1, collection1.getResourceId(), "2");
 
         List<ImmutablePair<PartitionKeyRange, IServerIdentity>> rangesBeforeSplit2 =
             new ArrayList<>();
@@ -141,7 +142,7 @@ public void before_AddressResolverTest() throws Exception {
         addPartitionKeyRangeFunc.apply(rangesBeforeSplit2);
 
 
-        this.routingMapCollection2BeforeSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesBeforeSplit2, collection2.getResourceId());
+        this.routingMapCollection2BeforeSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesBeforeSplit2, collection2.getResourceId(), "2");
 
         List<ImmutablePair<PartitionKeyRange, IServerIdentity>> rangesAfterSplit2 =
             new ArrayList<>();
@@ -162,7 +163,7 @@ public void before_AddressResolverTest() throws Exception {
         addPartitionKeyRangeFunc.apply(rangesAfterSplit2);
 
 
-        this.routingMapCollection2AfterSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesAfterSplit2, collection2.getResourceId());
+        this.routingMapCollection2AfterSplit = InMemoryCollectionRoutingMap.tryCreateCompleteRoutingMap(rangesAfterSplit2, collection2.getResourceId(), "2");
     }
 
     private void TestCacheRefreshWhileRouteByPartitionKey(

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/query/FeedResponseBuilder.java

@@ -3,7 +3,6 @@
 
 package com.azure.cosmos.implementation.query;
 
-import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.implementation.Resource;