Update cosmos emulator tests to include vnext

eng/pipelines/templates/stages/cosmos-emulator-matrix.json

@@ -11,13 +11,15 @@
       "Emulator Only Integration Tests Tcp - Java 8": {
           "ProfileFlag": "-Pemulator",
           "PROTOCOLS": "[\"Tcp\"]",
+          "CONNECTION_MODES": "[\"Direct\", \"Gateway\"]",
           "DESIRED_CONSISTENCIES": "[\"Session\"]",
           "JavaTestVersion": "1.8",
           "AdditionalArgs": "-DACCOUNT_HOST=https://localhost:8081/ -DCOSMOS.AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY=true"
       },
       "Emulator Only Integration Tests Tcp - Java 17": {
         "ProfileFlag": "-Pemulator",
         "PROTOCOLS": "[\"Tcp\"]",
+        "CONNECTION_MODES": "[\"Direct\", \"Gateway\"]",
         "DESIRED_CONSISTENCIES": "[\"Strong\"]",
         "JavaTestVersion": "1.17",
         "AdditionalArgs": "-DACCOUNT_HOST=https://localhost:8081/ -DCOSMOS.AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY=true"

eng/pipelines/templates/stages/cosmos-emulator-vnext-matrix.json

@@ -8,7 +8,7 @@
     },
     "EmulatorConfig": {
       "Emulator VNext Integration Tests On Insecure Connection - Java ": {
-        "ProfileFlag": "-Pemulator-vnext",
+        "ProfileFlag": "-Pemulator",
         "DESIRED_CONSISTENCIES": "[\"Session\"]",
         "AdditionalArgs": "-DCOSMOS.AZURE_COSMOS_DISABLE_NON_STREAMING_ORDER_BY=true"
       }

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

@@ -37,7 +37,7 @@ public class ChangeFeedContinuationTokenUtilsTests extends TestSuiteBase {
     private CosmosAsyncClient client;
     private CosmosAsyncDatabase createdDatabase;
 
-    @Factory(dataProvider = "simpleClientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public ChangeFeedContinuationTokenUtilsTests(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -156,7 +156,7 @@ public static Object[][] changeFeedSplitHandlingDataProvider() {
         };
     }
 
-    @Factory(dataProvider = "simpleClientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public CosmosContainerChangeFeedTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -28,7 +28,7 @@ public class CosmosContainerContentResponseOnWriteTest extends TestSuiteBase {
     private CosmosDatabase createdDatabase;
 
     //  Currently Gateway and Direct TCP support minimal response feature.
-    @Factory(dataProvider = "clientBuildersWithDirectTcpWithContentResponseOnWriteDisabled")
+    @Factory(dataProvider = "emulatorClientBuildersContentResponseOnWriteEnabledFalse")
     public CosmosContainerContentResponseOnWriteTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -28,7 +28,7 @@ public class CosmosDatabaseContentResponseOnWriteTest extends TestSuiteBase {
     private CosmosDatabase createdDatabase;
 
     //  Currently Gateway and Direct TCP support minimal response feature.
-    @Factory(dataProvider = "clientBuildersWithDirectTcpWithContentResponseOnWriteDisabled")
+    @Factory(dataProvider = "emulatorClientBuildersContentResponseOnWriteEnabledFalse")
     public CosmosDatabaseContentResponseOnWriteTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -54,7 +54,7 @@ public class CosmosDiagnosticsE2ETest extends TestSuiteBase {
 
     private CosmosClient client;
 
-    @Factory(dataProvider = "clientBuildersWithDirectSession")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public CosmosDiagnosticsE2ETest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -27,26 +27,26 @@ public class CosmosItemContentResponseOnWriteTest extends TestSuiteBase {
     private CosmosContainer container;
 
     //  Currently Gateway and Direct TCP support minimal response feature.
-    @Factory(dataProvider = "clientBuildersWithContentResponseOnWriteEnabledAndDisabled")
+    @Factory(dataProvider = "emulatorClientBuildersContentResponseOnWriteEnabledAndDisabled")
     public CosmosItemContentResponseOnWriteTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }
 
-    @BeforeClass(groups = {"fast"}, timeOut = SETUP_TIMEOUT)
+    @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT)
     public void beforeClass() {
         assertThat(this.client).isNull();
         this.client = getClientBuilder().buildClient();
         CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.client.asyncClient());
         container = client.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId());
     }
 
-    @AfterClass(groups = {"fast"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
+    @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)
     public void afterClass() {
         assertThat(this.client).isNotNull();
         this.client.close();
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void createItem_withContentResponseOnWriteDisabled() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemRequestOptions cosmosItemRequestOptions = new CosmosItemRequestOptions();
@@ -62,7 +62,7 @@ public void createItem_withContentResponseOnWriteDisabled() throws Exception {
         validateMinimalItemResponse(properties, itemResponse1, true);
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void createItem_withContentResponseOnWriteEnabledThroughRequestOptions() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemRequestOptions cosmosItemRequestOptions = new CosmosItemRequestOptions();
@@ -78,7 +78,7 @@ public void createItem_withContentResponseOnWriteEnabledThroughRequestOptions()
         validateItemResponse(properties, itemResponse1);
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void readItem_withContentResponseOnWriteDisabled() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemRequestOptions cosmosItemRequestOptions = new CosmosItemRequestOptions();
@@ -96,7 +96,7 @@ public void readItem_withContentResponseOnWriteDisabled() throws Exception {
 
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void readItem_withContentResponseOnWriteEnabledThroughRequestOptions() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemResponse<InternalObjectNode> itemResponse = container.createItem(properties);
@@ -115,7 +115,7 @@ public void readItem_withContentResponseOnWriteEnabledThroughRequestOptions() th
 
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void replaceItem_withContentResponseOnWriteDisabled() {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemRequestOptions cosmosItemRequestOptions = new CosmosItemRequestOptions();
@@ -137,7 +137,7 @@ public void replaceItem_withContentResponseOnWriteDisabled() {
         validateMinimalItemResponse(properties, replace, true);
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void replaceItem_withContentResponseOnWriteEnabledThroughRequestOptions() throws Exception{
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemRequestOptions cosmosItemRequestOptions = new CosmosItemRequestOptions();
@@ -159,7 +159,7 @@ public void replaceItem_withContentResponseOnWriteEnabledThroughRequestOptions()
         validateItemResponse(properties, replace);
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void deleteItem_withContentResponseOnWriteDisabled() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemResponse<InternalObjectNode> itemResponse = container.createItem(properties);
@@ -172,7 +172,7 @@ public void deleteItem_withContentResponseOnWriteDisabled() throws Exception {
         validateMinimalItemResponse(properties, deleteResponse, false);
     }
 
-    @Test(groups = { "fast" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator" }, timeOut = TIMEOUT)
     public void deleteItem_withContentResponseOnWriteEnabledThroughRequestOptions() throws Exception {
         InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString());
         CosmosItemResponse<InternalObjectNode> itemResponse = container.createItem(properties);

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

@@ -65,7 +65,7 @@ public class CosmosItemWriteRetriesTest extends TestSuiteBase {
 
     private TracerUnderTest mockTracer;
 
-    @Factory(dataProvider = "clientBuildersWithDirectTcpSession")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public CosmosItemWriteRetriesTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilderWithReducedNetworkRequestTimeout(clientBuilder));
     }

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

@@ -26,7 +26,7 @@ public class DistributedClientTest extends TestSuiteBase {
     private CosmosAsyncClient client;
     private CosmosAsyncContainer container;
 
-    @Factory(dataProvider = "clientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public DistributedClientTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -35,12 +35,12 @@ public class EmulatorVNextWithHttpTest extends TestSuiteBase {
     private CosmosAsyncContainer createdContainer;
     private CosmosAsyncDatabase createdDatabase;
 
-    @Factory(dataProvider = "clientBuilders")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public EmulatorVNextWithHttpTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }
 
-    @BeforeClass(groups = { "emulator-vnext" }, timeOut = SETUP_TIMEOUT)
+    @BeforeClass(groups = { "emulator", "emulator-vnext" }, timeOut = SETUP_TIMEOUT)
     public void before_EmulatorWithHttpTest() {
         this.client =
             getClientBuilder()
@@ -51,13 +51,13 @@ public void before_EmulatorWithHttpTest() {
         this.createdContainer = getSharedMultiPartitionCosmosContainerWithIdAsPartitionKey(this.client);
     }
 
-    @AfterClass(groups = { "emulator-vnext" }, timeOut = 3 * SHUTDOWN_TIMEOUT, alwaysRun = true)
+    @AfterClass(groups = { "emulator", "emulator-vnext" }, timeOut = 3 * SHUTDOWN_TIMEOUT, alwaysRun = true)
     public void afterClass() {
         logger.info("starting ....");
         safeClose(this.client);
     }
 
-    @Test(groups = { "emulator-vnext" }, timeOut = TIMEOUT)
+    @Test(groups = { "emulator", "emulator-vnext" }, timeOut = TIMEOUT)
     public void createSameDatabaseTwice() {
         try {
             this.client.createDatabase(this.createdDatabase.getId()).block();
@@ -69,7 +69,7 @@ public void createSameDatabaseTwice() {
         }
     }
 
-    @Test(groups = { "emulator-vnext" }, timeOut = 4 * TIMEOUT)
+    @Test(groups = { "emulator", "emulator-vnext" }, timeOut = 4 * TIMEOUT)
     public void documentCrud() {
         // Currently emulator vnext only support gateway mode and limited set of features
         // https://review.learn.microsoft.com/en-us/azure/cosmos-db/emulator-linux?branch=release-ignite-2024-cosmos-db#feature-support

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

@@ -29,7 +29,7 @@
 @Ignore("MalformedResponseTests is only safe to run in isolation as it leverages Reflection to override the ObjectMapper instance responsible for deserialization.")
 public class MalformedResponseTests extends TestSuiteBase {
 
-    @Factory(dataProvider = "clientBuildersWithSessionConsistency")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public MalformedResponseTests(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -40,7 +40,7 @@ public class MetadataThrottlingRetryPolicyTest extends TestSuiteBase {
     private CosmosAsyncClient client;
     private CosmosAsyncContainer container;
 
-    @Factory(dataProvider = "simpleClientBuildersWithoutRetryOnThrottledRequests")
+    @Factory(dataProvider = "emulatorClientBuildersWithoutRetryOnThrottledRequests")
     public MetadataThrottlingRetryPolicyTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
         this.subscriberValidationTimeout = TIMEOUT;

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

@@ -3,6 +3,7 @@
 package com.azure.cosmos.implementation;
 
 import com.azure.cosmos.ConsistencyLevel;
+import com.azure.cosmos.ConnectionMode;
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosNettyLeakDetectorFactory;
@@ -36,7 +37,6 @@
 import org.apache.commons.lang3.StringUtils;
 import org.mockito.Mockito;
 import org.mockito.stubbing.Answer;
-import org.testng.ITestContext;
 import org.testng.annotations.AfterSuite;
 import org.testng.annotations.BeforeSuite;
 import org.testng.annotations.DataProvider;
@@ -72,6 +72,7 @@ public abstract class TestSuiteBase extends DocumentClientTest {
     protected static final ImmutableList<String> preferredLocations;
     private static final ImmutableList<ConsistencyLevel> desiredConsistencies;
     private static final ImmutableList<Protocol> protocols;
+    private static final ImmutableList<ConnectionMode> connectionModes;
 
     protected int subscriberValidationTimeout = TIMEOUT;
     protected static Database SHARED_DATABASE;
@@ -92,6 +93,9 @@ private static <T> ImmutableList<T> immutableListOrNull(List<T> list) {
         preferredLocations = immutableListOrNull(parsePreferredLocation(TestConfigurations.PREFERRED_LOCATIONS));
         protocols = ObjectUtils.defaultIfNull(immutableListOrNull(parseProtocols(TestConfigurations.PROTOCOLS)),
                                               ImmutableList.of(Protocol.TCP));
+        // Defaulting to Gateway if no connection mode is specified to maintain backward compatibility
+        connectionModes = ObjectUtils.defaultIfNull(immutableListOrNull(parseConnectionModes(TestConfigurations.CONNECTION_MODES)),
+            ImmutableList.of(ConnectionMode.GATEWAY));
         //  Object mapper configuration
         objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
         objectMapper.configure(JsonParser.Feature.ALLOW_SINGLE_QUOTES, true);
@@ -830,6 +834,53 @@ public static <T extends Resource> void validateQueryFailure(Flux<FeedResponse<T
         validator.validate((Throwable) testSubscriber.getEvents().get(1).get(0));
     }
 
+    /**
+     * Returns CosmosClientBuilders for the static field connectionModes with the specified contentResponseOnWriteEnabled flag.
+     */
+    private static Object[][] emulatorClientBuildersWithContentResponseOnWriteEnabled(boolean contentResponseOnWriteEnabled) {
+        List<Object[]> builders = new ArrayList<>();
+        for (ConnectionMode mode : connectionModes) {
+            switch (mode) {
+                case DIRECT:
+                    builders.add(new Object[]{createDirectRxDocumentClient(
+                        ConsistencyLevel.SESSION,
+                        Protocol.TCP,
+                        false,
+                        preferredLocations,
+                        contentResponseOnWriteEnabled
+                    )});
+                    break;
+                case GATEWAY:
+                    builders.add(new Object[]{createGatewayRxDocumentClient(
+                        ConsistencyLevel.SESSION,
+                        false,
+                        preferredLocations,
+                        contentResponseOnWriteEnabled
+                    )});
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unsupported ConnectionMode: " + mode);
+            }
+        }
+        return builders.toArray(new Object[0][]);
+    }
+
+    /**
+     * DataProvider that returns CosmosClientBuilders for the static field connectionModes with contentResponseOnWriteEnabled = true.
+     */
+    @DataProvider
+    public static Object[][] emulatorClientBuilders() {
+        return emulatorClientBuildersWithContentResponseOnWriteEnabled(true);
+    }
+
+    /**
+     * DataProvider that returns CosmosClientBuilders for the static field connectionModes with contentResponseOnWriteEnabled = false.
+     */
+    @DataProvider
+    public static Object[][] emulatorClientBuildersContentResponseOnWriteEnabledFalse() {
+        return emulatorClientBuildersWithContentResponseOnWriteEnabled(false);
+    }
+
     @DataProvider
     public static Object[][] clientBuilders() {
         return new Object[][]{{createGatewayRxDocumentClient(ConsistencyLevel.SESSION, false, null, true)}};
@@ -886,6 +937,23 @@ static List<Protocol> parseProtocols(String protocols) {
         }
     }
 
+    static List<ConnectionMode> parseConnectionModes(String connectionModes) {
+        if (StringUtils.isEmpty(connectionModes)) {
+            return null;
+        }
+        List<ConnectionMode> modeList = new ArrayList<>();
+        try {
+            List<String> modeStrings = objectMapper.readValue(connectionModes, new com.fasterxml.jackson.core.type.TypeReference<List<String>>() {});
+            for (String mode : modeStrings) {
+                modeList.add(ConnectionMode.valueOf(com.azure.cosmos.implementation.guava25.base.CaseFormat.UPPER_CAMEL.to(com.azure.cosmos.implementation.guava25.base.CaseFormat.UPPER_UNDERSCORE, mode)));
+            }
+            return modeList;
+        } catch (Exception e) {
+            logger.error("INVALID configured test connectionModes [{}].", connectionModes);
+            throw new IllegalStateException("INVALID configured test connectionModes " + connectionModes);
+        }
+    }
+
     @DataProvider
     public static Object[][] simpleClientBuildersWithDirect() {
         return simpleClientBuildersWithDirect(true, toArray(protocols));

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

@@ -64,7 +64,7 @@ public class ThroughputControlTests extends TestSuiteBase {
     private CosmosAsyncDatabase database;
     private CosmosAsyncContainer container;
 
-    @Factory(dataProvider = "simpleClientBuildersWithoutRetryOnThrottledRequests")
+    @Factory(dataProvider = "emulatorClientBuildersWithoutRetryOnThrottledRequests")
     public ThroughputControlTests(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
         this.subscriberValidationTimeout = TIMEOUT;

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

@@ -26,7 +26,7 @@ public class ThroughputControlGroupConfigConfigurationTests extends TestSuiteBas
     private CosmosAsyncDatabase database;
     private CosmosAsyncContainer container;
 
-    @Factory(dataProvider = "clientBuildersWithSessionConsistency")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public ThroughputControlGroupConfigConfigurationTests(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
         this.subscriberValidationTimeout = TIMEOUT;

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

@@ -52,7 +52,7 @@ public class CollectionCrudTest extends TestSuiteBase {
     private CosmosAsyncClient client;
     private CosmosAsyncDatabase database;
 
-    @Factory(dataProvider = "clientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public CollectionCrudTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
         this.subscriberValidationTimeout = TIMEOUT;

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

@@ -25,7 +25,7 @@ public class DocumentClientResourceLeakTest extends TestSuiteBase {
     private CosmosAsyncDatabase createdDatabase;
     private CosmosAsyncContainer createdCollection;
 
-    @Factory(dataProvider = "simpleClientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public DocumentClientResourceLeakTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }

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

@@ -43,7 +43,7 @@ public class DocumentCrudTest extends TestSuiteBase {
     private CosmosAsyncContainer container;
     private CosmosAsyncDatabase database;
 
-    @Factory(dataProvider = "clientBuildersWithDirect")
+    @Factory(dataProvider = "emulatorClientBuilders")
     public DocumentCrudTest(CosmosClientBuilder clientBuilder) {
         super(clientBuilder);
     }