[cosmos] Use fixed self-owned accounts for key-based live tests

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosNotFoundTests.java

@@ -229,7 +229,7 @@ public void performDocumentOperationOnNonExistentContainerGatewayModeV2(Operatio
         }
     }
 
-    @Test(groups = {"thinclient"}, timeOut = TIMEOUT)
+    @Test(groups = {"thinclient"}, timeOut = 2 * TIMEOUT, retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void performBulkOnNonExistentContainerGatewayModeV2() {
         logger.info("Running test: Read item from non-existent container in Gateway Connection Mode");
 
@@ -306,7 +306,7 @@ public void performDocumentOperationOnDeletedContainer(OperationType operationTy
 
             CosmosAsyncContainer testContainer = clientToUse.getDatabase(testAsyncDatabase.getId()).getContainer(testContainerId);
 
-            Thread.sleep(5000);
+            Thread.sleep(15000);
 
             // Create a different client instance to delete the container
             deletingAsyncClient = getClientBuilder()
@@ -437,7 +437,8 @@ public void performBulkOnDeletedContainer() throws InterruptedException {
         }
     }
 
-    @Test(groups = {"thinclient"}, dataProvider = "operationTypeProvider", timeOut = TIMEOUT)
+    @Test(groups = {"thinclient"}, dataProvider = "operationTypeProvider",
+        timeOut = 2 * TIMEOUT, retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void performDocumentOperationOnDeletedContainerWithGatewayV2(OperationType operationType) throws InterruptedException {
         logger.info("Running test: Read item from deleted container - Gateway V2 Connection Mode");
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java

@@ -203,8 +203,10 @@ public void cosmosAsyncClient(
         createAndInitializeDiagnosticsProvider(
             mockTracer, useLegacyTracing, enableRequestLevelTracing, ShowQueryMode.NONE, forceThresholdViolations, samplingRate);
 
-        CosmosDatabaseResponse cosmosDatabaseResponse = client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId(),
-            ThroughputProperties.createManualThroughput(5000)).block();
+        CosmosDatabaseResponse cosmosDatabaseResponse = executeControlPlaneWithRetry(
+            () -> client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId(),
+                ThroughputProperties.createManualThroughput(5000)).block(),
+            mockTracer::reset);
         assertThat(cosmosDatabaseResponse).isNotNull();
         verifyTracerAttributes(
             mockTracer,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readAllAfterCreation.java

@@ -12,7 +12,8 @@
 import java.util.function.Function;
 
 public class FITests_readAllAfterCreation extends FaultInjectionWithAvailabilityStrategyTestsBase {
-    @Test(groups = {"fi-multi-master", "fi-thinclient-multi-master"}, dataProvider = "testConfigs_readAllAfterCreation", retryAnalyzer = FlakyTestRetryAnalyzer.class)
+    @Test(groups = {"fi-multi-master", "fi-thinclient-multi-master"},
+        dataProvider = "testConfigs_readAllAfterCreation", retryAnalyzer = SuperFlakyTestRetryAnalyzer.class)
     public void readAllAfterCreation(
         String testCaseId,
         Duration endToEndTimeout,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FITests_readManyByPartitionKeysAfterCreation.java

@@ -36,7 +36,8 @@
 public class FITests_readManyByPartitionKeysAfterCreation
     extends FaultInjectionWithAvailabilityStrategyTestsBase {
 
-    @Test(groups = {"fi-multi-master"}, dataProvider = "testConfigs_readManyByPartitionKeysAfterCreation", retryAnalyzer = FlakyTestRetryAnalyzer.class)
+    @Test(groups = {"fi-multi-master"}, dataProvider = "testConfigs_readManyByPartitionKeysAfterCreation",
+        retryAnalyzer = SuperFlakyTestRetryAnalyzer.class)
     public void readManyByPartitionKeysAfterCreation(
         String testCaseId,
         Duration endToEndTimeout,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/FaultInjectionWithAvailabilityStrategyTestsBase.java

@@ -5082,7 +5082,8 @@ private CosmosAsyncContainer createTestContainer(CosmosAsyncClient clientWithPre
     private CosmosAsyncContainer createTestContainer(CosmosAsyncClient clientWithPreferredRegions, String dbId) {
         String containerId = UUID.randomUUID().toString();
 
-        clientWithPreferredRegions.createDatabaseIfNotExists(dbId).block();
+        executeControlPlaneWithRetry(
+            () -> clientWithPreferredRegions.createDatabaseIfNotExists(dbId).block());
         CosmosAsyncDatabase databaseWithSeveralWriteableRegions = clientWithPreferredRegions.getDatabase(dbId);
 
         // setup db and container and pass their ids accordingly

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ProactiveConnectionManagementTest.java

@@ -256,7 +256,8 @@ public void openConnectionsAndInitCachesWithContainer(ProactiveConnectionManagem
         }
     }
 
-    @Test(groups = {"multi-master"}, dataProvider = "proactiveContainerInitConfigs")
+    @Test(groups = {"multi-master"}, dataProvider = "proactiveContainerInitConfigs",
+        retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void openConnectionsAndInitCachesWithCosmosClient_And_PerContainerConnectionPoolSize_ThroughSystemConfig(
         ProactiveConnectionManagementTestConfig proactiveConnectionManagementTestConfig) throws InterruptedException {
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SessionConsistencyWithRegionScopingTests.java

@@ -1806,7 +1806,8 @@ public void readYouWriteWithNoExplicitRegionSwitching(
         }
     }
 
-    @Test(groups = {"multi-master"}, dataProvider = "readYouWriteWithExplicitRegionSwitchingTestContext", timeOut = 80 * TIMEOUT)
+    @Test(groups = {"multi-master"}, dataProvider = "readYouWriteWithExplicitRegionSwitchingTestContext",
+        timeOut = 80 * TIMEOUT, retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void readYouWriteWithExplicitRegionSwitching(
         BiFunction<CosmosAsyncContainer, Boolean, Set<String>> func,
         String testId,

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/SplitTestsRetryAnalyzer.java

@@ -12,10 +12,28 @@ public SplitTestsRetryAnalyzer() {
 
     @Override
     public boolean retry(ITestResult iTestResult) {
-        if (!(iTestResult.getThrowable() instanceof SplitTimeoutException)) {
+        Throwable throwable = iTestResult.getThrowable();
+        // Forcing a partition split is inherently disruptive. Besides the explicit
+        // SplitTimeoutException, a split on a shared live-test account can surface transient
+        // 408 (timeout), 429 (throttling) and 503 (service unavailable) errors while the split
+        // is in progress. Retry all of these (bounded by retryLimit) so split tests are not
+        // flaky on transient contention; deterministic failures are still surfaced.
+        if (!(throwable instanceof SplitTimeoutException) && !isTransientCosmosFailure(throwable)) {
             return false;
         }
 
         return super.retry(iTestResult);
     }
+
+    private static boolean isTransientCosmosFailure(Throwable throwable) {
+        Throwable current = throwable;
+        while (current != null) {
+            if (current instanceof CosmosException) {
+                int statusCode = ((CosmosException) current).getStatusCode();
+                return statusCode == 408 || statusCode == 429 || statusCode == 503;
+            }
+            current = current.getCause();
+        }
+        return false;
+    }
 }

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/faultinjection/FaultInjectionMetadataRequestRuleTests.java

@@ -540,24 +540,31 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError(
         int applyLimit,
         boolean shouldInjectPreferredRegionsOnClient) throws JsonProcessingException {
 
-        // We need to create a new client because client may have marked region unavailable in other tests
-        // which can impact the test result
-        CosmosAsyncClient testClient = getClientBuilder()
-            .contentResponseOnWriteEnabled(true)
-            .preferredRegions(shouldInjectPreferredRegionsOnClient ? this.writePreferredLocations : Collections.emptyList())
-            .endToEndOperationLatencyPolicyConfig(
-                new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofMinutes(10)).build())
-            .buildAsyncClient();
+        String originalSharedPkRangeCacheSetting =
+            System.getProperty("COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED");
+        System.setProperty("COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED", "false");
 
-        CosmosAsyncContainer container =
-            testClient
-                .getDatabase(this.cosmosAsyncContainer.getDatabase().getId())
-                .getContainer(this.cosmosAsyncContainer.getId());
+        CosmosAsyncClient testClient = null;
+        FaultInjectionRule pkRangesConnectionDelayRule = null;
+        FaultInjectionRule dataOperationGoneRule = null;
+        try {
+            // We need to create a new client because client may have marked region unavailable in other tests
+            // which can impact the test result
+            testClient = getClientBuilder()
+                .contentResponseOnWriteEnabled(true)
+                .preferredRegions(shouldInjectPreferredRegionsOnClient ? this.writePreferredLocations : Collections.emptyList())
+                .endToEndOperationLatencyPolicyConfig(
+                    new CosmosEndToEndOperationLatencyPolicyConfigBuilder(Duration.ofMinutes(10)).build())
+                .buildAsyncClient();
 
-        // Test to validate partition key ranges request is being injected connection timeout
-        String pkRangesConnectionDelay = "PkRanges-connectionDelay-" + UUID.randomUUID();
-        FaultInjectionRule pkRangesConnectionDelayRule =
-            new FaultInjectionRuleBuilder(pkRangesConnectionDelay)
+            CosmosAsyncContainer container =
+                testClient
+                    .getDatabase(this.cosmosAsyncContainer.getDatabase().getId())
+                    .getContainer(this.cosmosAsyncContainer.getId());
+
+            // Test to validate partition key ranges request is being injected connection timeout
+            String pkRangesConnectionDelay = "PkRanges-connectionDelay-" + UUID.randomUUID();
+            pkRangesConnectionDelayRule = new FaultInjectionRuleBuilder(pkRangesConnectionDelay)
                 .condition(
                     new FaultInjectionConditionBuilder()
                         .region(this.writePreferredLocations.get(0))
@@ -574,8 +581,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError(
                 .duration(Duration.ofMinutes(5))
                 .build();
 
-        FaultInjectionRule dataOperationGoneRule =
-            new FaultInjectionRuleBuilder("DataOperation-gone-" + UUID.randomUUID())
+            dataOperationGoneRule = new FaultInjectionRuleBuilder("DataOperation-gone-" + UUID.randomUUID())
                 .condition(
                     new FaultInjectionConditionBuilder()
                         .operationType(FaultInjectionOperationType.CREATE_ITEM)
@@ -589,7 +595,7 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError(
                 )
                 .duration(Duration.ofMinutes(5))
                 .build();
-        try {
+
             // create few items to first make sure the collection cache, pkRanges cache is being populated
             for (int i = 0; i < 10; i++) {
                 container.createItem(TestObject.create()).block();
@@ -602,6 +608,8 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError(
 
             try {
                 CosmosDiagnostics cosmosDiagnostics = container.createItem(TestObject.create()).block().getDiagnostics();
+                long ruleHitCount = pkRangesConnectionDelayRule.getHitCount();
+                assertThat(ruleHitCount).isGreaterThan(0);
                 // The PkRanges requests may have retried in another region,
                 // but the create request will only be retried locally for PARTITION_IS_SPLITTING
                 assertThat(cosmosDiagnostics.getContactedRegionNames().size()).isEqualTo(1);
@@ -627,19 +635,31 @@ public void faultInjectionServerErrorRuleTests_PartitionKeyRanges_DelayError(
 
                 assertThat(pkRangesLookup).isNotNull();
                 if (faultInjectionServerErrorType == FaultInjectionServerErrorType.CONNECTION_DELAY) {
-                    assertThat(pkRangesLookup.get("durationinMS").asLong()).isGreaterThanOrEqualTo(45 * 1000 * Math.min(applyLimit, 3)); // the duration will be at least one connection timeout
+                    assertThat(pkRangesLookup.get("durationinMS").asLong())
+                        .isGreaterThanOrEqualTo(45 * 1000 * Math.min(ruleHitCount, 3));
                 }
 
                 if (faultInjectionServerErrorType == FaultInjectionServerErrorType.RESPONSE_DELAY) {
-                    assertThat(pkRangesLookup.get("durationinMS").asLong()).isGreaterThanOrEqualTo(500 * Math.min(applyLimit, 3)); // the duration will be at least one response timeout
+                    assertThat(pkRangesLookup.get("durationinMS").asLong())
+                        .isGreaterThanOrEqualTo(500 * Math.min(ruleHitCount, 3));
                 }
 
             } catch (CosmosException cosmosException) {
                 fail("CreateItem should have succeeded. " + cosmosException.getDiagnostics());
             }
         } finally {
-            pkRangesConnectionDelayRule.disable();
-            dataOperationGoneRule.disable();
+            if (originalSharedPkRangeCacheSetting == null) {
+                System.clearProperty("COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED");
+            } else {
+                System.setProperty(
+                    "COSMOS.SHARED_PARTITION_KEY_RANGE_CACHE_ENABLED", originalSharedPkRangeCacheSetting);
+            }
+            if (pkRangesConnectionDelayRule != null) {
+                pkRangesConnectionDelayRule.disable();
+            }
+            if (dataOperationGoneRule != null) {
+                dataOperationGoneRule.disable();
+            }
             safeClose(testClient);
         }
     }

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

@@ -8,6 +8,7 @@
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.DirectConnectionConfig;
+import com.azure.cosmos.FlakyTestRetryAnalyzer;
 import com.azure.cosmos.GatewayConnectionConfig;
 import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
 import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils;
@@ -143,7 +144,8 @@ public void validateEventualConsistencyOnAsyncReplicationGateway() {
         // https://msdata.visualstudio.com/CosmosDB/_workitems/edit/355053
     }
 
-    @Test(groups = {"direct"}, dataProvider = "regionScopedSessionContainerConfigs", timeOut = CONSISTENCY_TEST_TIMEOUT)
+    @Test(groups = {"direct"}, dataProvider = "regionScopedSessionContainerConfigs",
+        timeOut = CONSISTENCY_TEST_TIMEOUT, retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void validateSessionContainerAfterCollectionDeletion(boolean shouldRegionScopedSessionContainerEnabled) throws Exception {
         //TODO Need to test with TCP protocol
         // https://msdata.visualstudio.com/CosmosDB/_workitems/edit/355057

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

@@ -144,6 +144,30 @@ public void validate(List<FeedResponse<T>> feedList) {
             return this;
         }
 
+        /**
+         * Asserts the results CONTAIN all of the expected resource IDs, allowing additional
+         * unexpected results. Unlike {@link #exactlyContainsInAnyOrder(List)} this does not
+         * require the result set to be limited to the expected IDs, so it is safe for
+         * account-global reads (e.g. readAllDatabases) that run against shared test accounts
+         * where other resources may exist concurrently.
+         */
+        public Builder<T> containsResourceIds(List<String> expectedIds) {
+            validators.add(new FeedResponseListValidator<T>() {
+                @Override
+                public void validate(List<FeedResponse<T>> feedList) {
+                    List<String> actualIds = feedList
+                            .stream()
+                            .flatMap(f -> f.getResults().stream())
+                            .map(r -> getResource(r).getResourceId())
+                            .collect(Collectors.toList());
+                    assertThat(actualIds)
+                    .describedAs("Resource IDs of results")
+                    .containsAll(expectedIds);
+                }
+            });
+            return this;
+        }
+
         public Builder<T> exactlyContainsIdsInAnyOrder(List<String> expectedIds) {
             validators.add(new FeedResponseListValidator<T>() {
                 @Override

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

@@ -152,6 +152,7 @@ public void setThinClientHttpClient_triggersProbeOnRefresh() throws Exception {
 
             // Probe is fire-and-forget on a scheduler -> wait briefly for it to run.
             waitForProbeCallCount(probeCallCount, 2, Duration.ofSeconds(5));
+            waitForProxyDecision(gem, Boolean.TRUE, Duration.ofSeconds(5));
 
             assertThat(probeCallCount.get()).as("probe was issued for each thin-client region").isGreaterThanOrEqualTo(2);
             assertThat(gem.getProxyProbeDecision()).as("after all-200 cycle, proxy is healthy").isEqualTo(Boolean.TRUE);

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

@@ -1383,10 +1383,10 @@ public void tryGetAddress_repeatedlySetUnhealthyStatus_forceRefresh() throws Int
         // 2. replica validation will get triggered in case of unhealthyPending / unknown addresses, replica validation will do a
         // submitOpenConnectionTaskOutsideLoop for each of these addresses but before that it will also do
         // isCollectionRidUnderOpenConnectionsFlow check to determine the no. of connections to open
-        Mockito.verify(proactiveOpenConnectionsProcessorMock, Mockito.times(2))
+        Mockito.verify(proactiveOpenConnectionsProcessorMock, Mockito.timeout(5000).times(2))
                .isCollectionRidUnderOpenConnectionsFlow(Mockito.any());
 
-        Mockito.verify(proactiveOpenConnectionsProcessorMock, Mockito.times(1))
+        Mockito.verify(proactiveOpenConnectionsProcessorMock, Mockito.timeout(5000).times(1))
                .submitOpenConnectionTaskOutsideLoop(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.anyInt());
     }
 

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/ChangeFeedTest.java

@@ -5,6 +5,7 @@
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.CosmosDiagnostics;
+import com.azure.cosmos.FlakyTestRetryAnalyzer;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.ClientSideRequestStatistics;
 import com.azure.cosmos.implementation.Database;
@@ -387,7 +388,7 @@ public void changeFeed_fromStartDate() throws Exception {
         assertThat(count).as("Change feed should have one newly created document").isEqualTo(1);
     }
 
-    @Test(groups = { "query" }, timeOut = TIMEOUT)
+    @Test(groups = { "query" }, timeOut = TIMEOUT, retryAnalyzer = FlakyTestRetryAnalyzer.class)
     public void changesFromPartitionKey_AfterInsertingNewDocuments() throws Exception {
         String partitionKey = partitionKeyToDocuments.keySet().iterator().next();
         FeedRange feedRange = new FeedRangePartitionKeyImpl(

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/MultiMasterConflictResolutionTest.java

@@ -50,7 +50,9 @@ public void conflictResolutionPolicyCRUD() {
 
         // default last writer wins, path _ts
         CosmosContainerProperties containerSettings = new CosmosContainerProperties(UUID.randomUUID().toString(), partitionKeyDef);
-        database.createContainer(containerSettings, new CosmosContainerRequestOptions()).block();
+        final CosmosContainerProperties initialContainerSettings = containerSettings;
+        executeControlPlaneWithRetry(() ->
+            database.createContainer(initialContainerSettings, new CosmosContainerRequestOptions()).block());
         CosmosAsyncContainer container = database.getContainer(containerSettings.getId());
 
         containerSettings = container.read().block().getProperties();
@@ -124,7 +126,9 @@ private void testConflictResolutionPolicyRequiringPath(ConflictResolutionMode co
             } else {
                 collectionSettings.setConflictResolutionPolicy(ConflictResolutionPolicy.createCustomPolicy(paths[i]));
             }
-            collectionSettings = database.createContainer(collectionSettings, new CosmosContainerRequestOptions()).block().getProperties();
+            final CosmosContainerProperties containerToCreate = collectionSettings;
+            collectionSettings = executeControlPlaneWithRetry(() ->
+                database.createContainer(containerToCreate, new CosmosContainerRequestOptions()).block().getProperties());
             assertThat(collectionSettings.getConflictResolutionPolicy().getMode()).isEqualTo(conflictResolutionMode);
 
             if (conflictResolutionMode == ConflictResolutionMode.LAST_WRITER_WINS) {