Improving Test robustness after changing configured Regions for live tests

sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/KafkaCosmosConnectContainer.java

@@ -3,29 +3,41 @@
 
 package com.azure.cosmos.kafka.connect;
 
-import com.azure.core.exception.ResourceNotFoundException;
 import org.apache.kafka.clients.admin.AdminClient;
 import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.errors.TopicExistsException;
+import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.sourcelab.kafka.connect.apiclient.Configuration;
 import org.sourcelab.kafka.connect.apiclient.KafkaConnectClient;
 import org.sourcelab.kafka.connect.apiclient.request.dto.ConnectorDefinition;
 import org.sourcelab.kafka.connect.apiclient.request.dto.ConnectorStatus;
 import org.sourcelab.kafka.connect.apiclient.request.dto.NewConnectorDefinition;
+import org.sourcelab.kafka.connect.apiclient.rest.exceptions.InvalidRequestException;
+import org.sourcelab.kafka.connect.apiclient.rest.exceptions.ResourceNotFoundException;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.KafkaContainer;
+import org.testcontainers.containers.wait.strategy.Wait;
 import org.testcontainers.utility.DockerImageName;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.Properties;
 import java.util.UUID;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public class KafkaCosmosConnectContainer extends GenericContainer<KafkaCosmosConnectContainer> {
     private static final Logger logger = LoggerFactory.getLogger(KafkaCosmosConnectContainer.class);
     private static final int KAFKA_CONNECT_PORT = 8083;
+    private static final Duration KAFKA_CONNECT_REST_OPERATION_TIMEOUT = Duration.ofMinutes(2);
+    private static final Duration KAFKA_CONNECT_REST_RETRY_DELAY = Duration.ofMillis(500);
+    private static final int KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS = 30;
     private Properties producerProperties;
     private Properties consumerProperties;
     private AdminClient adminClient;
@@ -54,6 +66,10 @@ private void defaultConfig() {
 //        withEnv("CONNECT_LOG4J_LOGGERS", "org.apache.kafka=DEBUG,org.reflections=DEBUG,com.azure.cosmos.kafka=DEBUG");
 
         withExposedPorts(KAFKA_CONNECT_PORT);
+        waitingFor(Wait.forHttp("/connectors")
+            .forPort(KAFKA_CONNECT_PORT)
+            .forStatusCode(200)
+            .withStartupTimeout(KAFKA_CONNECT_REST_OPERATION_TIMEOUT));
     }
 
     private Properties defaultConsumerConfig() {
@@ -158,13 +174,9 @@ public void registerConnector(String name, Map<String, String> config) {
         KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
 
         logger.info("adding kafka connector {}", name);
-
-        try {
-            Thread.sleep(500);
-        } catch (InterruptedException e) {
-            throw new RuntimeException(e);
-        }
-        ConnectorDefinition connectorDefinition = kafkaConnectClient.addConnector(newConnectorDefinition);
+        ConnectorDefinition connectorDefinition = executeWithKafkaConnectRestRetry(
+            "adding kafka connector " + name,
+            () -> kafkaConnectClient.addConnector(newConnectorDefinition));
         logger.info("adding kafka connector completed with " + connectorDefinition);
     }
 
@@ -212,7 +224,9 @@ public void resumeConnector(String name) {
 
     public ConnectorStatus getConnectorStatus(String name) {
         KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
-        return kafkaConnectClient.getConnectorStatus(name);
+        return executeWithKafkaConnectRestRetry(
+            "getting kafka connector status " + name,
+            () -> kafkaConnectClient.getConnectorStatus(name));
     }
 
     public String getTarget() {
@@ -232,11 +246,91 @@ public Properties getConsumerProperties() {
     }
 
     public void createTopic(String topicName, int numPartitions) {
-        this.adminClient.createTopics(
-            Arrays.asList(new NewTopic(topicName, numPartitions, (short) replicationFactor)));
+        try {
+            this.adminClient.createTopics(
+                    Arrays.asList(new NewTopic(topicName, numPartitions, (short) replicationFactor)))
+                .all()
+                .get(KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
+            logger.info("Creating topic {} succeeded.", topicName);
+        } catch (ExecutionException exception) {
+            if (exception.getCause() instanceof TopicExistsException) {
+                logger.info("Topic {} already exists.", topicName);
+                return;
+            }
+
+            throw new RuntimeException("Failed to create topic " + topicName, exception);
+        } catch (InterruptedException exception) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException("Interrupted while creating topic " + topicName, exception);
+        } catch (TimeoutException exception) {
+            throw new RuntimeException("Timed out while creating topic " + topicName, exception);
+        }
     }
 
     public void deleteTopic(String topicName) {
-        this.adminClient.deleteTopics(Arrays.asList(topicName));
+        try {
+            this.adminClient.deleteTopics(Arrays.asList(topicName))
+                .all()
+                .get(KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
+            logger.info("Deleting topic {} succeeded.", topicName);
+        } catch (ExecutionException exception) {
+            if (exception.getCause() instanceof UnknownTopicOrPartitionException) {
+                logger.info("Topic {} not found.", topicName);
+                return;
+            }
+
+            logger.warn("Failed to delete topic {}", topicName, exception);
+        } catch (InterruptedException exception) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException("Interrupted while deleting topic " + topicName, exception);
+        } catch (TimeoutException exception) {
+            logger.warn("Timed out while deleting topic {}", topicName, exception);
+        }
+    }
+
+    private <T> T executeWithKafkaConnectRestRetry(String operationName, Callable<T> operation) {
+        long deadlineNanos = System.nanoTime() + KAFKA_CONNECT_REST_OPERATION_TIMEOUT.toNanos();
+        int attempts = 0;
+        InvalidRequestException lastException = null;
+
+        while (System.nanoTime() < deadlineNanos) {
+            attempts++;
+            try {
+                return operation.call();
+            } catch (InvalidRequestException exception) {
+                if (!isTransientKafkaConnectRestNotFound(exception)) {
+                    throw exception;
+                }
+
+                lastException = exception;
+                logger.warn(
+                    "Kafka Connect REST returned transient Not Found while {} on attempt {}. Retrying.",
+                    operationName,
+                    attempts,
+                    exception);
+            } catch (Exception exception) {
+                throw new RuntimeException("Failed while " + operationName, exception);
+            }
+
+            sleepBeforeKafkaConnectRestRetry(operationName);
+        }
+
+        throw new RuntimeException(
+            "Timed out after " + KAFKA_CONNECT_REST_OPERATION_TIMEOUT.getSeconds()
+                + " seconds while " + operationName,
+            lastException);
+    }
+
+    private static boolean isTransientKafkaConnectRestNotFound(InvalidRequestException exception) {
+        return exception.getErrorCode() == 404;
+    }
+
+    private static void sleepBeforeKafkaConnectRestRetry(String operationName) {
+        try {
+            TimeUnit.MILLISECONDS.sleep(KAFKA_CONNECT_REST_RETRY_DELAY.toMillis());
+        } catch (InterruptedException exception) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException("Interrupted while " + operationName, exception);
+        }
     }
 }

sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/KafkaCosmosTestSuiteBase.java

@@ -19,6 +19,8 @@
 import com.azure.cosmos.models.IncludedPath;
 import com.azure.cosmos.models.IndexingPolicy;
 import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.azure.cosmos.models.SqlParameter;
+import com.azure.cosmos.models.SqlQuerySpec;
 import com.azure.cosmos.models.ThroughputProperties;
 import com.fasterxml.jackson.databind.JsonNode;
 import org.apache.commons.lang3.StringUtils;
@@ -35,9 +37,11 @@
 import java.lang.reflect.Method;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.UUID;
+import java.util.concurrent.TimeUnit;
 
 @Listeners({KafkaCosmosTestNGLogListener.class})
 public class KafkaCosmosTestSuiteBase implements ITest {
@@ -46,6 +50,9 @@ public class KafkaCosmosTestSuiteBase implements ITest {
 
     protected static final int SUITE_SETUP_TIMEOUT = 120000;
     protected static final int SUITE_SHUTDOWN_TIMEOUT = 60000;
+    private static final int KAFKA_COSMOS_SUITE_SETUP_TIMEOUT = 10 * SUITE_SETUP_TIMEOUT;
+    private static final Duration CONTAINER_METADATA_MAX_WAIT = Duration.ofMinutes(2);
+    private static final Duration CONTAINER_METADATA_ATTEMPT_TIMEOUT = Duration.ofSeconds(10);
 
     protected static final AzureKeyCredential credential;
     protected static String databaseName;
@@ -89,7 +96,7 @@ protected static CosmosContainerProperties getSinglePartitionContainer(CosmosAsy
         credential = new AzureKeyCredential(KafkaCosmosTestConfigurations.MASTER_KEY);
     }
 
-    @BeforeSuite(groups = { "kafka", "kafka-integration" }, timeOut = SUITE_SETUP_TIMEOUT)
+    @BeforeSuite(groups = { "kafka", "kafka-integration" }, timeOut = KAFKA_COSMOS_SUITE_SETUP_TIMEOUT)
     public void beforeSuite() {
 
         logger.info("beforeSuite Started");
@@ -119,9 +126,11 @@ public void beforeSuite() {
                     options,
                     6000);
         }
+
+        waitForCreatedContainersToBeQueryable();
     }
 
-    @BeforeSuite(groups = { "kafka-emulator" }, timeOut = SUITE_SETUP_TIMEOUT)
+    @BeforeSuite(groups = { "kafka-emulator" }, timeOut = KAFKA_COSMOS_SUITE_SETUP_TIMEOUT)
     public void beforeSuite_emulator() {
 
         logger.info("beforeSuite Started");
@@ -151,6 +160,8 @@ public void beforeSuite_emulator() {
                     options,
                     6000);
         }
+
+        waitForCreatedContainersToBeQueryable();
     }
 
     @BeforeSuite(groups = { "unit" }, timeOut = SUITE_SETUP_TIMEOUT)
@@ -227,6 +238,98 @@ private static String createCollection(
         return cosmosContainerProperties.getId();
     }
 
+    private static void waitForCreatedContainersToBeQueryable() {
+        try (CosmosAsyncClient probeClient = createGatewayHouseKeepingDocumentClient(true).buildAsyncClient()) {
+            waitForCreatedContainersToBeQueryable(
+                probeClient,
+                databaseName,
+                Arrays.asList(
+                    multiPartitionContainerName,
+                    multiPartitionContainerWithIdAsPartitionKeyName,
+                    singlePartitionContainerName));
+        }
+    }
+
+    private static void waitForCreatedContainersToBeQueryable(
+        CosmosAsyncClient cosmosAsyncClient,
+        String databaseName,
+        List<String> expectedContainerNames) {
+
+        long deadlineNanos = System.nanoTime() + CONTAINER_METADATA_MAX_WAIT.toNanos();
+        int attempts = 0;
+        Throwable lastFailure = null;
+
+        while (System.nanoTime() < deadlineNanos) {
+            attempts++;
+            try {
+                List<String> visibleContainerNames = getVisibleContainerNames(
+                    cosmosAsyncClient,
+                    databaseName,
+                    expectedContainerNames);
+
+                if (visibleContainerNames.containsAll(expectedContainerNames)) {
+                    logger.info(
+                        "Kafka test containers {} became queryable in database {} after {} attempt(s).",
+                        expectedContainerNames,
+                        databaseName,
+                        attempts);
+                    return;
+                }
+
+                lastFailure = new AssertionError(
+                    "Expected containers " + expectedContainerNames + " but only found " + visibleContainerNames);
+            } catch (Exception exception) {
+                lastFailure = exception;
+            }
+
+            try {
+                TimeUnit.MILLISECONDS.sleep(500);
+            } catch (InterruptedException exception) {
+                Thread.currentThread().interrupt();
+                throw new RuntimeException("Interrupted while waiting for Kafka test containers to become queryable.", exception);
+            }
+        }
+
+        throw new AssertionError(
+            "Kafka test containers " + expectedContainerNames + " were not queryable in database "
+                + databaseName + " within " + CONTAINER_METADATA_MAX_WAIT.getSeconds() + " seconds after "
+                + attempts + " attempt(s).",
+            lastFailure);
+    }
+
+    private static List<String> getVisibleContainerNames(
+        CosmosAsyncClient cosmosAsyncClient,
+        String databaseName,
+        List<String> expectedContainerNames) {
+
+        StringBuilder queryBuilder = new StringBuilder("SELECT * FROM c WHERE c.id IN (");
+        List<SqlParameter> parameters = new ArrayList<>();
+        for (int index = 0; index < expectedContainerNames.size(); index++) {
+            String parameterName = "@container" + index;
+            parameters.add(new SqlParameter(parameterName, expectedContainerNames.get(index)));
+            queryBuilder.append(parameterName);
+            if (index < expectedContainerNames.size() - 1) {
+                queryBuilder.append(", ");
+            }
+        }
+        queryBuilder.append(")");
+
+        List<CosmosContainerProperties> visibleContainers = cosmosAsyncClient
+            .getDatabase(databaseName)
+            .queryContainers(new SqlQuerySpec(queryBuilder.toString(), parameters))
+            .byPage()
+            .flatMapIterable(response -> response.getResults())
+            .collectList()
+            .block(CONTAINER_METADATA_ATTEMPT_TIMEOUT);
+
+        List<String> visibleContainerNames = new ArrayList<>();
+        for (CosmosContainerProperties visibleContainer : visibleContainers) {
+            visibleContainerNames.add(visibleContainer.getId());
+        }
+
+        return visibleContainerNames;
+    }
+
     static protected CosmosContainerProperties getCollectionDefinitionWithRangeRangeIndex(boolean enableAllVersionsAndDeletesPolicy) {
         return getCollectionDefinitionWithRangeRangeIndex(Collections.singletonList("/mypk"), enableAllVersionsAndDeletesPolicy);
     }

sdk/cosmos/azure-cosmos-spark_3-3_2-12/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Added a defensive guard in bounded change feed reads (with `endLsn`) that fails the Spark task with `IllegalStateException` when the underlying paginator stops before the latest continuation token has advanced to `endLsn`. - See [PR 49393](https://github.com/Azure/azure-sdk-for-java/pull/49393)
+* Fixed an issue in the `readManyByPartitionKeys` API in the Spark connector which could result in duplicates and missing the first record. - See [PR 49694](https://github.com/Azure/azure-sdk-for-java/pull/49694)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3-4_2-12/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Added a defensive guard in bounded change feed reads (with `endLsn`) that fails the Spark task with `IllegalStateException` when the underlying paginator stops before the latest continuation token has advanced to `endLsn`. - See [PR 49393](https://github.com/Azure/azure-sdk-for-java/pull/49393)
+* Fixed an issue in the `readManyByPartitionKeys` API in the Spark connector which could result in duplicates and missing the first record. - See [PR 49694](https://github.com/Azure/azure-sdk-for-java/pull/49694)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3-5_2-12/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Added a defensive guard in bounded change feed reads (with `endLsn`) that fails the Spark task with `IllegalStateException` when the underlying paginator stops before the latest continuation token has advanced to `endLsn`. - See [PR 49393](https://github.com/Azure/azure-sdk-for-java/pull/49393)
+* Fixed an issue in the `readManyByPartitionKeys` API in the Spark connector which could result in duplicates and missing the first record. - See [PR 49694](https://github.com/Azure/azure-sdk-for-java/pull/49694)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos-spark_3-5_2-13/CHANGELOG.md

@@ -8,6 +8,7 @@
 
 #### Bugs Fixed
 * Added a defensive guard in bounded change feed reads (with `endLsn`) that fails the Spark task with `IllegalStateException` when the underlying paginator stops before the latest continuation token has advanced to `endLsn`. - See [PR 49393](https://github.com/Azure/azure-sdk-for-java/pull/49393)
+* Fixed an issue in the `readManyByPartitionKeys` API in the Spark connector which could result in duplicates and missing the first record. - See [PR 49694](https://github.com/Azure/azure-sdk-for-java/pull/49694)
 
 #### Other Changes