Skip to content

Commit a79d582

Browse files
FabianMeiswinkelCopilot
FabianMeiswinkel
and
Copilot
authored
Improving Test robustness after changing configured Regions for live tests (#49639)
* Add more verbose asserts * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Fix NITs in tests * Update PerPartitionCircuitBreakerE2ETests.java * Update ReadFeedStoredProceduresTest.java * Update CosmosItemTest.java * Update StoredProcedureUpsertReplaceTest.java * Update FaultInjectionServerErrorRuleOnDirectTests.java * Update FaultInjectionServerErrorRuleOnDirectTests.java * Update StoredProcedureQueryTest.java * Fix SP tests * More SP test fixes * Update CosmosContainerChangeFeedTest.java * Fix test flakiness * Update CosmosSyncStoredProcTest.java * Fix test flakiness * Update TestSuiteBase.java * Changing test regions * Update TestSuiteBase.java * Update CosmosDiagnosticsTest.java * Update KafkaCosmosConnectContainer.java * Fix test issues * Update TestSuiteBase.java * Update ClientMetricsTest.java * Update ClientMetricsTest.java * Update EndToEndTimeOutValidationTests.java * Update CosmosItemTest.java * Fix test flakiness * Address flakiness * Systemically fixing container creation race * Update TestSuiteBase.java * Update TransactionalBatchAsyncContainerTest.java * Update EndToEndTimeOutValidationTests.java * Update EndToEndTimeOutValidationTests.java * Update EndToEndTimeOutValidationTests.java * Update CosmosItemTest.java * Update CosmosBulkTest.java * Update AzureKeyCredentialTest.java * Update CosmosContainerOpenConnectionsAndInitCachesTest.java * Update OperationPoliciesTest.java * Update ProactiveConnectionManagementTest.java * Fixed flaky tests * Update CosmosItemContentResponseOnWriteTest.java * Make getFeedRanges from tests more robust * Update GatewayAddressCacheTest.java * Update CustomerWorkflowPartitionLevelCircuitBreakerTest.java * Fix Test flakiness * Update CosmosItemTest.java * Increasing setup timeouts * More robustness imprvements * More robustness fixes * Update LocationCacheTest.java * Update FaultInjectionWithAvailabilityStrategyTestsBase.java * Update TestSuiteBase.java * Fix FI unit tests * Fixing test flakiness * Update TestSuiteBase.java * fix bulk test failures * Update FaultInjectionWithAvailabilityStrategyTestsBase.java * Update TestSuiteBase.java * Update OperationPoliciesTest.java * Fixes ReadManyByPK bug in Spark connector resulting in duplicates and possibly missing first record * Updating changelogs * Update CosmosReadManyByPartitionKeyReader.scala * Reacting to code review feedback * Fixing remaining test failures * Fixing remaining test flakiness --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent 932ddbe commit a79d582

78 files changed

Lines changed: 3402 additions & 914 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 106 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,29 +3,41 @@
33

44
package com.azure.cosmos.kafka.connect;
55

6-
import com.azure.core.exception.ResourceNotFoundException;
76
import org.apache.kafka.clients.admin.AdminClient;
87
import org.apache.kafka.clients.admin.NewTopic;
98
import org.apache.kafka.clients.producer.ProducerConfig;
9+
import org.apache.kafka.common.errors.TopicExistsException;
10+
import org.apache.kafka.common.errors.UnknownTopicOrPartitionException;
1011
import org.slf4j.Logger;
1112
import org.slf4j.LoggerFactory;
1213
import org.sourcelab.kafka.connect.apiclient.Configuration;
1314
import org.sourcelab.kafka.connect.apiclient.KafkaConnectClient;
1415
import org.sourcelab.kafka.connect.apiclient.request.dto.ConnectorDefinition;
1516
import org.sourcelab.kafka.connect.apiclient.request.dto.ConnectorStatus;
1617
import org.sourcelab.kafka.connect.apiclient.request.dto.NewConnectorDefinition;
18+
import org.sourcelab.kafka.connect.apiclient.rest.exceptions.InvalidRequestException;
19+
import org.sourcelab.kafka.connect.apiclient.rest.exceptions.ResourceNotFoundException;
1720
import org.testcontainers.containers.GenericContainer;
1821
import org.testcontainers.containers.KafkaContainer;
22+
import org.testcontainers.containers.wait.strategy.Wait;
1923
import org.testcontainers.utility.DockerImageName;
2024

25+
import java.time.Duration;
2126
import java.util.Arrays;
2227
import java.util.Map;
2328
import java.util.Properties;
2429
import java.util.UUID;
30+
import java.util.concurrent.Callable;
31+
import java.util.concurrent.ExecutionException;
32+
import java.util.concurrent.TimeUnit;
33+
import java.util.concurrent.TimeoutException;
2534

2635
public class KafkaCosmosConnectContainer extends GenericContainer<KafkaCosmosConnectContainer> {
2736
private static final Logger logger = LoggerFactory.getLogger(KafkaCosmosConnectContainer.class);
2837
private static final int KAFKA_CONNECT_PORT = 8083;
38+
private static final Duration KAFKA_CONNECT_REST_OPERATION_TIMEOUT = Duration.ofMinutes(2);
39+
private static final Duration KAFKA_CONNECT_REST_RETRY_DELAY = Duration.ofMillis(500);
40+
private static final int KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS = 30;
2941
private Properties producerProperties;
3042
private Properties consumerProperties;
3143
private AdminClient adminClient;
@@ -54,6 +66,10 @@ private void defaultConfig() {
5466
// withEnv("CONNECT_LOG4J_LOGGERS", "org.apache.kafka=DEBUG,org.reflections=DEBUG,com.azure.cosmos.kafka=DEBUG");
5567

5668
withExposedPorts(KAFKA_CONNECT_PORT);
69+
waitingFor(Wait.forHttp("/connectors")
70+
.forPort(KAFKA_CONNECT_PORT)
71+
.forStatusCode(200)
72+
.withStartupTimeout(KAFKA_CONNECT_REST_OPERATION_TIMEOUT));
5773
}
5874

5975
private Properties defaultConsumerConfig() {
@@ -158,13 +174,9 @@ public void registerConnector(String name, Map<String, String> config) {
158174
KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
159175

160176
logger.info("adding kafka connector {}", name);
161-
162-
try {
163-
Thread.sleep(500);
164-
} catch (InterruptedException e) {
165-
throw new RuntimeException(e);
166-
}
167-
ConnectorDefinition connectorDefinition = kafkaConnectClient.addConnector(newConnectorDefinition);
177+
ConnectorDefinition connectorDefinition = executeWithKafkaConnectRestRetry(
178+
"adding kafka connector " + name,
179+
() -> kafkaConnectClient.addConnector(newConnectorDefinition));
168180
logger.info("adding kafka connector completed with " + connectorDefinition);
169181
}
170182

@@ -212,7 +224,9 @@ public void resumeConnector(String name) {
212224

213225
public ConnectorStatus getConnectorStatus(String name) {
214226
KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
215-
return kafkaConnectClient.getConnectorStatus(name);
227+
return executeWithKafkaConnectRestRetry(
228+
"getting kafka connector status " + name,
229+
() -> kafkaConnectClient.getConnectorStatus(name));
216230
}
217231

218232
public String getTarget() {
@@ -232,11 +246,91 @@ public Properties getConsumerProperties() {
232246
}
233247

234248
public void createTopic(String topicName, int numPartitions) {
235-
this.adminClient.createTopics(
236-
Arrays.asList(new NewTopic(topicName, numPartitions, (short) replicationFactor)));
249+
try {
250+
this.adminClient.createTopics(
251+
Arrays.asList(new NewTopic(topicName, numPartitions, (short) replicationFactor)))
252+
.all()
253+
.get(KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
254+
logger.info("Creating topic {} succeeded.", topicName);
255+
} catch (ExecutionException exception) {
256+
if (exception.getCause() instanceof TopicExistsException) {
257+
logger.info("Topic {} already exists.", topicName);
258+
return;
259+
}
260+
261+
throw new RuntimeException("Failed to create topic " + topicName, exception);
262+
} catch (InterruptedException exception) {
263+
Thread.currentThread().interrupt();
264+
throw new RuntimeException("Interrupted while creating topic " + topicName, exception);
265+
} catch (TimeoutException exception) {
266+
throw new RuntimeException("Timed out while creating topic " + topicName, exception);
267+
}
237268
}
238269

239270
public void deleteTopic(String topicName) {
240-
this.adminClient.deleteTopics(Arrays.asList(topicName));
271+
try {
272+
this.adminClient.deleteTopics(Arrays.asList(topicName))
273+
.all()
274+
.get(KAFKA_ADMIN_OPERATION_TIMEOUT_IN_SECONDS, TimeUnit.SECONDS);
275+
logger.info("Deleting topic {} succeeded.", topicName);
276+
} catch (ExecutionException exception) {
277+
if (exception.getCause() instanceof UnknownTopicOrPartitionException) {
278+
logger.info("Topic {} not found.", topicName);
279+
return;
280+
}
281+
282+
logger.warn("Failed to delete topic {}", topicName, exception);
283+
} catch (InterruptedException exception) {
284+
Thread.currentThread().interrupt();
285+
throw new RuntimeException("Interrupted while deleting topic " + topicName, exception);
286+
} catch (TimeoutException exception) {
287+
logger.warn("Timed out while deleting topic {}", topicName, exception);
288+
}
289+
}
290+
291+
private <T> T executeWithKafkaConnectRestRetry(String operationName, Callable<T> operation) {
292+
long deadlineNanos = System.nanoTime() + KAFKA_CONNECT_REST_OPERATION_TIMEOUT.toNanos();
293+
int attempts = 0;
294+
InvalidRequestException lastException = null;
295+
296+
while (System.nanoTime() < deadlineNanos) {
297+
attempts++;
298+
try {
299+
return operation.call();
300+
} catch (InvalidRequestException exception) {
301+
if (!isTransientKafkaConnectRestNotFound(exception)) {
302+
throw exception;
303+
}
304+
305+
lastException = exception;
306+
logger.warn(
307+
"Kafka Connect REST returned transient Not Found while {} on attempt {}. Retrying.",
308+
operationName,
309+
attempts,
310+
exception);
311+
} catch (Exception exception) {
312+
throw new RuntimeException("Failed while " + operationName, exception);
313+
}
314+
315+
sleepBeforeKafkaConnectRestRetry(operationName);
316+
}
317+
318+
throw new RuntimeException(
319+
"Timed out after " + KAFKA_CONNECT_REST_OPERATION_TIMEOUT.getSeconds()
320+
+ " seconds while " + operationName,
321+
lastException);
322+
}
323+
324+
private static boolean isTransientKafkaConnectRestNotFound(InvalidRequestException exception) {
325+
return exception.getErrorCode() == 404;
326+
}
327+
328+
private static void sleepBeforeKafkaConnectRestRetry(String operationName) {
329+
try {
330+
TimeUnit.MILLISECONDS.sleep(KAFKA_CONNECT_REST_RETRY_DELAY.toMillis());
331+
} catch (InterruptedException exception) {
332+
Thread.currentThread().interrupt();
333+
throw new RuntimeException("Interrupted while " + operationName, exception);
334+
}
241335
}
242336
}

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

Lines changed: 105 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import com.azure.cosmos.models.IncludedPath;
2020
import com.azure.cosmos.models.IndexingPolicy;
2121
import com.azure.cosmos.models.PartitionKeyDefinition;
22+
import com.azure.cosmos.models.SqlParameter;
23+
import com.azure.cosmos.models.SqlQuerySpec;
2224
import com.azure.cosmos.models.ThroughputProperties;
2325
import com.fasterxml.jackson.databind.JsonNode;
2426
import org.apache.commons.lang3.StringUtils;
@@ -35,9 +37,11 @@
3537
import java.lang.reflect.Method;
3638
import java.time.Duration;
3739
import java.util.ArrayList;
40+
import java.util.Arrays;
3841
import java.util.Collections;
3942
import java.util.List;
4043
import java.util.UUID;
44+
import java.util.concurrent.TimeUnit;
4145

4246
@Listeners({KafkaCosmosTestNGLogListener.class})
4347
public class KafkaCosmosTestSuiteBase implements ITest {
@@ -46,6 +50,9 @@ public class KafkaCosmosTestSuiteBase implements ITest {
4650

4751
protected static final int SUITE_SETUP_TIMEOUT = 120000;
4852
protected static final int SUITE_SHUTDOWN_TIMEOUT = 60000;
53+
private static final int KAFKA_COSMOS_SUITE_SETUP_TIMEOUT = 10 * SUITE_SETUP_TIMEOUT;
54+
private static final Duration CONTAINER_METADATA_MAX_WAIT = Duration.ofMinutes(2);
55+
private static final Duration CONTAINER_METADATA_ATTEMPT_TIMEOUT = Duration.ofSeconds(10);
4956

5057
protected static final AzureKeyCredential credential;
5158
protected static String databaseName;
@@ -89,7 +96,7 @@ protected static CosmosContainerProperties getSinglePartitionContainer(CosmosAsy
8996
credential = new AzureKeyCredential(KafkaCosmosTestConfigurations.MASTER_KEY);
9097
}
9198

92-
@BeforeSuite(groups = { "kafka", "kafka-integration" }, timeOut = SUITE_SETUP_TIMEOUT)
99+
@BeforeSuite(groups = { "kafka", "kafka-integration" }, timeOut = KAFKA_COSMOS_SUITE_SETUP_TIMEOUT)
93100
public void beforeSuite() {
94101

95102
logger.info("beforeSuite Started");
@@ -119,9 +126,11 @@ public void beforeSuite() {
119126
options,
120127
6000);
121128
}
129+
130+
waitForCreatedContainersToBeQueryable();
122131
}
123132

124-
@BeforeSuite(groups = { "kafka-emulator" }, timeOut = SUITE_SETUP_TIMEOUT)
133+
@BeforeSuite(groups = { "kafka-emulator" }, timeOut = KAFKA_COSMOS_SUITE_SETUP_TIMEOUT)
125134
public void beforeSuite_emulator() {
126135

127136
logger.info("beforeSuite Started");
@@ -151,6 +160,8 @@ public void beforeSuite_emulator() {
151160
options,
152161
6000);
153162
}
163+
164+
waitForCreatedContainersToBeQueryable();
154165
}
155166

156167
@BeforeSuite(groups = { "unit" }, timeOut = SUITE_SETUP_TIMEOUT)
@@ -227,6 +238,98 @@ private static String createCollection(
227238
return cosmosContainerProperties.getId();
228239
}
229240

241+
private static void waitForCreatedContainersToBeQueryable() {
242+
try (CosmosAsyncClient probeClient = createGatewayHouseKeepingDocumentClient(true).buildAsyncClient()) {
243+
waitForCreatedContainersToBeQueryable(
244+
probeClient,
245+
databaseName,
246+
Arrays.asList(
247+
multiPartitionContainerName,
248+
multiPartitionContainerWithIdAsPartitionKeyName,
249+
singlePartitionContainerName));
250+
}
251+
}
252+
253+
private static void waitForCreatedContainersToBeQueryable(
254+
CosmosAsyncClient cosmosAsyncClient,
255+
String databaseName,
256+
List<String> expectedContainerNames) {
257+
258+
long deadlineNanos = System.nanoTime() + CONTAINER_METADATA_MAX_WAIT.toNanos();
259+
int attempts = 0;
260+
Throwable lastFailure = null;
261+
262+
while (System.nanoTime() < deadlineNanos) {
263+
attempts++;
264+
try {
265+
List<String> visibleContainerNames = getVisibleContainerNames(
266+
cosmosAsyncClient,
267+
databaseName,
268+
expectedContainerNames);
269+
270+
if (visibleContainerNames.containsAll(expectedContainerNames)) {
271+
logger.info(
272+
"Kafka test containers {} became queryable in database {} after {} attempt(s).",
273+
expectedContainerNames,
274+
databaseName,
275+
attempts);
276+
return;
277+
}
278+
279+
lastFailure = new AssertionError(
280+
"Expected containers " + expectedContainerNames + " but only found " + visibleContainerNames);
281+
} catch (Exception exception) {
282+
lastFailure = exception;
283+
}
284+
285+
try {
286+
TimeUnit.MILLISECONDS.sleep(500);
287+
} catch (InterruptedException exception) {
288+
Thread.currentThread().interrupt();
289+
throw new RuntimeException("Interrupted while waiting for Kafka test containers to become queryable.", exception);
290+
}
291+
}
292+
293+
throw new AssertionError(
294+
"Kafka test containers " + expectedContainerNames + " were not queryable in database "
295+
+ databaseName + " within " + CONTAINER_METADATA_MAX_WAIT.getSeconds() + " seconds after "
296+
+ attempts + " attempt(s).",
297+
lastFailure);
298+
}
299+
300+
private static List<String> getVisibleContainerNames(
301+
CosmosAsyncClient cosmosAsyncClient,
302+
String databaseName,
303+
List<String> expectedContainerNames) {
304+
305+
StringBuilder queryBuilder = new StringBuilder("SELECT * FROM c WHERE c.id IN (");
306+
List<SqlParameter> parameters = new ArrayList<>();
307+
for (int index = 0; index < expectedContainerNames.size(); index++) {
308+
String parameterName = "@container" + index;
309+
parameters.add(new SqlParameter(parameterName, expectedContainerNames.get(index)));
310+
queryBuilder.append(parameterName);
311+
if (index < expectedContainerNames.size() - 1) {
312+
queryBuilder.append(", ");
313+
}
314+
}
315+
queryBuilder.append(")");
316+
317+
List<CosmosContainerProperties> visibleContainers = cosmosAsyncClient
318+
.getDatabase(databaseName)
319+
.queryContainers(new SqlQuerySpec(queryBuilder.toString(), parameters))
320+
.byPage()
321+
.flatMapIterable(response -> response.getResults())
322+
.collectList()
323+
.block(CONTAINER_METADATA_ATTEMPT_TIMEOUT);
324+
325+
List<String> visibleContainerNames = new ArrayList<>();
326+
for (CosmosContainerProperties visibleContainer : visibleContainers) {
327+
visibleContainerNames.add(visibleContainer.getId());
328+
}
329+
330+
return visibleContainerNames;
331+
}
332+
230333
static protected CosmosContainerProperties getCollectionDefinitionWithRangeRangeIndex(boolean enableAllVersionsAndDeletesPolicy) {
231334
return getCollectionDefinitionWithRangeRangeIndex(Collections.singletonList("/mypk"), enableAllVersionsAndDeletesPolicy);
232335
}

‎sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionRuleBuilder.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -164,17 +164,27 @@ private void validateRuleOnGatewayConnection() {
164164
throw new IllegalArgumentException("STALED_ADDRESSES exception can not be injected for rule with gateway connection type");
165165
}
166166

167-
// for metadata request related rule, only CONNECTION_DELAY, RESPONSE_DELAY, TOO_MANY_REQUEST error can be injected
167+
// for metadata request related rule, only metadata-safe errors can be injected
168168
if (ImplementationBridgeHelpers
169169
.FaultInjectionConditionHelper
170170
.getFaultInjectionConditionAccessor()
171171
.isMetadataOperationType(this.condition)) {
172-
if (serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.TOO_MANY_REQUEST
173-
&& serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.RESPONSE_DELAY
174-
&& serverErrorResult.getServerErrorType() != FaultInjectionServerErrorType.CONNECTION_DELAY) {
172+
if (!isSupportedMetadataServerErrorType(serverErrorResult.getServerErrorType())) {
175173

176174
throw new IllegalArgumentException("Error type " + serverErrorResult.getServerErrorType() + " is not supported for rule with metadata request");
177175
}
178176
}
179177
}
178+
179+
private boolean isSupportedMetadataServerErrorType(FaultInjectionServerErrorType serverErrorType) {
180+
if (serverErrorType == FaultInjectionServerErrorType.TOO_MANY_REQUEST
181+
|| serverErrorType == FaultInjectionServerErrorType.RESPONSE_DELAY
182+
|| serverErrorType == FaultInjectionServerErrorType.CONNECTION_DELAY) {
183+
return true;
184+
}
185+
186+
return this.condition.getOperationType() == FaultInjectionOperationType.METADATA_REQUEST_PARTITION_KEY_RANGES
187+
&& (serverErrorType == FaultInjectionServerErrorType.OWNER_RESOURCE_NOT_EXISTS
188+
|| serverErrorType == FaultInjectionServerErrorType.COLLECTION_NOT_AVAILABLE_FOR_READ);
189+
}
180190
}

‎sdk/cosmos/azure-cosmos-test/src/main/java/com/azure/cosmos/test/faultinjection/FaultInjectionServerErrorType.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,12 @@ public enum FaultInjectionServerErrorType {
2323
/** 404-1002 from server */
2424
READ_SESSION_NOT_AVAILABLE,
2525

26+
/** 404-1003 from server */
27+
OWNER_RESOURCE_NOT_EXISTS,
28+
29+
/** 404-1013 from server */
30+
COLLECTION_NOT_AVAILABLE_FOR_READ,
31+
2632
/** 408 from server */
2733
TIMEOUT,
2834

0 commit comments

Comments
 (0)