multi-region
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java
index 19bd22d6f570..2761269269c0 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/PerPartitionAutomaticFailoverE2ETests.java
@@ -276,8 +276,11 @@ public class PerPartitionAutomaticFailoverE2ETests extends TestSuiteBase {
if (expectedResponseCharacteristics.shouldFinalResponseHaveSuccess) {
assertThat(cosmosDiagnostics.getDiagnosticsContext()).isNotNull();
- assertThat(cosmosDiagnostics.getDiagnosticsContext().getStatusCode() >= HttpConstants.StatusCodes.OK
- && cosmosDiagnostics.getDiagnosticsContext().getStatusCode() <= HttpConstants.StatusCodes.NOT_MODIFIED).isTrue();
+ int finalStatusCode = cosmosDiagnostics.getDiagnosticsContext().getStatusCode();
+ assertThat(finalStatusCode)
+ .as("final response status code should indicate success (2xx/304) but was %d (sub-status %d)",
+ finalStatusCode, cosmosDiagnostics.getDiagnosticsContext().getSubStatusCode())
+ .isBetween(HttpConstants.StatusCodes.OK, HttpConstants.StatusCodes.NOT_MODIFIED);
}
};
@@ -1818,7 +1821,7 @@ public void testPpafWithWriteFailoverWithEligibleErrorStatusCodesWithPpafDynamic
* Dynamic enablement is achieved by overriding GlobalEndpointManager's owner to
* inject the PPAF flag into DatabaseAccount snapshots.
*/
- @Test(groups = {"multi-region", "fi-thinclient-multi-region"}, dataProvider = "ppafNonWriteDynamicEnablementScenarios")
+ @Test(groups = {"multi-region", "fi-thinclient-multi-region"}, dataProvider = "ppafNonWriteDynamicEnablementScenarios", retryAnalyzer = FlakyTestRetryAnalyzer.class)
public void testFailoverBehaviorForNonWriteOperationsWithPpafDynamicEnablement(
String testType,
OperationType operationType,
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/SessionTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/SessionTest.java
index fc7cbdde87a5..f0fe5fbb64f2 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/SessionTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/SessionTest.java
@@ -5,6 +5,7 @@
import com.azure.cosmos.ConnectionMode;
import com.azure.cosmos.ConsistencyLevel;
+import com.azure.cosmos.FlakyTestRetryAnalyzer;
import com.azure.cosmos.ReadConsistencyStrategy;
import com.azure.cosmos.implementation.batch.ItemBatchOperation;
import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest;
@@ -338,7 +339,7 @@ public void partitionedSessionToken(boolean isNameBased) throws NoSuchMethodExce
safeClose(dummyState);
}
- @Test(groups = { "fast" }, timeOut = TIMEOUT, dataProvider = "sessionTestArgProvider")
+ @Test(groups = { "fast" }, timeOut = TIMEOUT, dataProvider = "sessionTestArgProvider", retryAnalyzer = FlakyTestRetryAnalyzer.class)
public void sessionTokenNotRequired(boolean isNameBased) {
spyClient.readCollection(getCollectionLink(isNameBased), null).block();
// No session token set for the master resource related request
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ContainerCreateDeleteWithSameNameTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ContainerCreateDeleteWithSameNameTest.java
index a6423d69bb03..e1013fe00a06 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ContainerCreateDeleteWithSameNameTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ContainerCreateDeleteWithSameNameTest.java
@@ -889,6 +889,9 @@ private void createDeleteContainerWithSameName(
String pkPathAfterRecreate,
Function getPkAfterRecreate) throws InterruptedException {
CosmosAsyncContainer container = null;
+ // A throwaway client runs the post-create readiness probe so it does not warm this test's main client
+ // collection cache - the test relies on that cache being stale after the container is recreated.
+ CosmosAsyncClient probeClient = getClientBuilder().buildAsyncClient();
try {
// step1: create container
String testContainerId = UUID.randomUUID().toString();
@@ -899,7 +902,7 @@ private void createDeleteContainerWithSameName(
partitionKeyDef.setPaths(paths);
CosmosContainerProperties containerProperties = getCollectionDefinition(testContainerId, partitionKeyDef);
- container = createCollection(this.createdDatabase, containerProperties, new CosmosContainerRequestOptions(), ruBeforeDelete);
+ container = createCollection(this.createdDatabase, containerProperties, new CosmosContainerRequestOptions(), ruBeforeDelete, probeClient);
// Step2: execute func
validateFunc.accept(container, getPkBeforeDelete, false);
@@ -912,13 +915,14 @@ private void createDeleteContainerWithSameName(
partitionKeyDef.setPaths(Arrays.asList(pkPathAfterRecreate));
containerProperties = getCollectionDefinition(testContainerId, partitionKeyDef);
- container = createCollection(this.createdDatabase, containerProperties, new CosmosContainerRequestOptions(), ruAfterRecreate);
+ container = createCollection(this.createdDatabase, containerProperties, new CosmosContainerRequestOptions(), ruAfterRecreate, probeClient);
// step5: same as step2.
// This part will confirm the cache refreshed correctly
validateFunc.accept(container, getPkAfterRecreate, true);
} finally {
safeDeleteCollection(container);
+ safeClose(probeClient);
}
}
@@ -930,6 +934,9 @@ private void changeFeedCreateDeleteContainerWithSameName(
String pkPathAfterRecreate) throws InterruptedException {
CosmosAsyncContainer feedContainer = null;
CosmosAsyncContainer leaseContainer = null;
+ // A throwaway client runs the post-create readiness probe so it does not warm this test's main client
+ // collection cache - the test relies on that cache being stale after the feed container is recreated.
+ CosmosAsyncClient probeClient = getClientBuilder().buildAsyncClient();
try {
// step1: create feed container and lease container
@@ -937,7 +944,7 @@ private void changeFeedCreateDeleteContainerWithSameName(
PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition();
partitionKeyDefinition.setPaths(Arrays.asList(pkPathBeforeDelete));
CosmosContainerProperties feedContainerProperties = getCollectionDefinition(feedContainerId, partitionKeyDefinition);
- feedContainer = createCollection(this.createdDatabase, feedContainerProperties, new CosmosContainerRequestOptions(), ruBeforeDelete);
+ feedContainer = createCollection(this.createdDatabase, feedContainerProperties, new CosmosContainerRequestOptions(), ruBeforeDelete, probeClient);
String leaseContainerId = UUID.randomUUID().toString();
CosmosContainerProperties leaseContainerProperties = getCollectionDefinition(leaseContainerId);
@@ -954,7 +961,7 @@ private void changeFeedCreateDeleteContainerWithSameName(
// step 4: recreate the feed container with same id as step 1
partitionKeyDefinition.setPaths(Arrays.asList(pkPathAfterRecreate));
feedContainerProperties = getCollectionDefinition(feedContainerId, partitionKeyDefinition);
- feedContainer = createCollection(this.createdDatabase, feedContainerProperties, new CosmosContainerRequestOptions(), ruAfterRecreate);
+ feedContainer = createCollection(this.createdDatabase, feedContainerProperties, new CosmosContainerRequestOptions(), ruAfterRecreate, probeClient);
// step5: recreate the lease container and lease container with same ids as step1
leaseContainer = createLeaseContainer(leaseContainerProperties.getId());
@@ -965,6 +972,7 @@ private void changeFeedCreateDeleteContainerWithSameName(
} finally {
safeDeleteCollection(feedContainer);
safeDeleteCollection(leaseContainer);
+ safeClose(probeClient);
}
}
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java
index 9a6427952ca6..4166dc386944 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiOrderByQueryTests.java
@@ -120,6 +120,7 @@ public void before_MultiOrderByQueryTests() throws Exception {
documents = new ArrayList<>();
client = getClientBuilder().buildAsyncClient();
documentCollection = getSharedMultiPartitionCosmosContainerWithCompositeAndSpatialIndexes(client);
+ waitForCollectionToBeAvailableToRead(documentCollection, /* probeClient */ null);
cleanUpContainer(documentCollection);
expectCount(documentCollection, 0);
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java
index c632fd06e315..bcfc6cb66bf4 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java
@@ -102,7 +102,7 @@ public void before_NonStreamingOrderByQueryVectorSearchTest() {
database.createContainer(containerProperties).block();
largeDataContainer = database.getContainer(largeDataContainerId);
- waitForCollectionToBeAvailableToRead();
+ waitForCollectionToBeAvailableToRead(largeDataContainer, /* probeClient */ null);
for (Document doc : getVectorDocs()) {
flatIndexContainer.createItem(doc).block();
diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java
index 16722fd466bd..1bf138fc529e 100644
--- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java
+++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/OrderbyDocumentQueryTest.java
@@ -688,7 +688,7 @@ public void before_OrderbyDocumentQueryTest() throws Exception {
}))
.block();
roundTripsContainer = createdDatabase.getContainer(containerName);
- waitForCollectionToBeAvailableToRead();
+ waitForCollectionToBeAvailableToRead(roundTripsContainer, /* probeClient */ null);
setupRoundTripContainer();
List