Skip to content

Commit bb5ba74

Browse files
FabianMeiswinkelFabianMeiswinkel
authored
Test changes (#47233)
1 parent 1e4bb36 commit bb5ba74

File tree

12 files changed

+234
-223
lines changed

12 files changed

+234
-223
lines changed

‎sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/ReadMyWriteWorkflow.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -498,4 +498,13 @@ protected String getDocumentLink(Document doc) {
498498
return doc.getSelfLink();
499499
}
500500
}
501+
502+
@Override
503+
void shutdown() {
504+
if (this.client != null) {
505+
this.client.close();
506+
}
507+
508+
super.shutdown();
509+
}
501510
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,7 @@ public void validateNoChargeOnFailedSessionRead() throws Exception {
233233
new CosmosClientTelemetryConfig()
234234
.sendClientTelemetryToService(ClientTelemetry.DEFAULT_CLIENT_TELEMETRY_ENABLED))
235235
.build();
236+
236237
QueryFeedOperationState dummyState = null;
237238
try {
238239
// CREATE collection

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

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,7 @@ public void queryWithContinuationTokenLimit(CosmosQueryRequestOptions options, S
9595

9696
client.clearCapturedRequests();
9797

98-
QueryFeedOperationState dummyState = TestUtils
99-
.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
98+
QueryFeedOperationState dummyState = TestUtils.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
10099
try {
101100
Flux<FeedResponse<Document>> queryObservable = client
102101
.queryDocuments(
@@ -148,6 +147,11 @@ public Document createDocument(AsyncDocumentClient client, String collectionLink
148147
@BeforeClass(groups = { "fast" }, timeOut = SETUP_TIMEOUT)
149148
public void before_DocumentQuerySpyWireContentTest() throws Exception {
150149

150+
SpyClientUnderTestFactory.ClientUnderTest oldSnapshot = client;
151+
if (oldSnapshot != null) {
152+
oldSnapshot.close();
153+
}
154+
151155
client = new SpyClientBuilder(this.clientBuilder()).build();
152156

153157
createdDatabase = SHARED_DATABASE;
@@ -177,18 +181,14 @@ public void before_DocumentQuerySpyWireContentTest() throws Exception {
177181
options,
178182
client
179183
);
180-
try {
181184

182-
// do the query once to ensure the collection is cached.
183-
client.queryDocuments(getMultiPartitionCollectionLink(), "select * from root", state, Document.class)
184-
.then().block();
185+
// do the query once to ensure the collection is cached.
186+
client.queryDocuments(getMultiPartitionCollectionLink(), "select * from root", state, Document.class)
187+
.then().block();
185188

186-
// do the query once to ensure the collection is cached.
187-
client.queryDocuments(getSinglePartitionCollectionLink(), "select * from root", state, Document.class)
188-
.then().block();
189-
} finally {
190-
safeClose(state);
191-
}
189+
// do the query once to ensure the collection is cached.
190+
client.queryDocuments(getSinglePartitionCollectionLink(), "select * from root", state, Document.class)
191+
.then().block();
192192
}
193193

194194
@AfterClass(groups = { "fast" }, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true)

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -137,8 +137,7 @@ public void queryWithMaxIntegratedCacheStaleness(CosmosQueryRequestOptions optio
137137

138138
client.clearCapturedRequests();
139139

140-
QueryFeedOperationState dummyState = TestUtils
141-
.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
140+
QueryFeedOperationState dummyState = TestUtils.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
142141
try {
143142
client.queryDocuments(
144143
collectionLink,
@@ -175,7 +174,6 @@ public void queryWithMaxIntegratedCacheStalenessInNanoseconds() {
175174
);
176175

177176
try {
178-
179177
assertThatThrownBy(() -> client
180178
.queryDocuments(collectionLink, query, state, Document.class)
181179
.blockLast())

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

Lines changed: 110 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -203,132 +203,129 @@ public void partitionedSessionToken(boolean isNameBased) throws NoSuchMethodExce
203203
spyClient
204204
);
205205

206-
try {
206+
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
207+
assertThat(getSessionTokensInRequests()).hasSize(1);
208+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
209+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
207210

208-
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
209-
assertThat(getSessionTokensInRequests()).hasSize(1);
210-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
211-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
211+
// Session token validation for cross partition query
212+
spyClient.clearCapturedRequests();
213+
queryRequestOptions = new CosmosQueryRequestOptions();
212214

213-
// Session token validation for cross partition query
214-
spyClient.clearCapturedRequests();
215-
queryRequestOptions = new CosmosQueryRequestOptions();
216-
safeClose(dummyState);
217-
dummyState = TestUtils.createDummyQueryFeedOperationState(
218-
ResourceType.Document,
219-
OperationType.Query,
220-
queryRequestOptions,
221-
spyClient
222-
);
223-
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
224-
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
225-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
226-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
215+
safeClose(dummyState);
216+
dummyState = TestUtils.createDummyQueryFeedOperationState(
217+
ResourceType.Document,
218+
OperationType.Query,
219+
queryRequestOptions,
220+
spyClient
221+
);
222+
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
223+
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
224+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
225+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
227226

228-
// Session token validation for feed ranges query
229-
spyClient.clearCapturedRequests();
230-
List<FeedRange> feedRanges = spyClient.getFeedRanges(getCollectionLink(isNameBased), true).block();
231-
queryRequestOptions = new CosmosQueryRequestOptions();
232-
queryRequestOptions.setFeedRange(feedRanges.get(0));
233-
safeClose(dummyState);
234-
dummyState = TestUtils.createDummyQueryFeedOperationState(
235-
ResourceType.Document,
236-
OperationType.Query,
237-
queryRequestOptions,
238-
spyClient
239-
);
240-
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
241-
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
242-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
243-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
227+
// Session token validation for feed ranges query
228+
spyClient.clearCapturedRequests();
229+
List<FeedRange> feedRanges = spyClient.getFeedRanges(getCollectionLink(isNameBased), true).block();
230+
queryRequestOptions = new CosmosQueryRequestOptions();
231+
queryRequestOptions.setFeedRange(feedRanges.get(0));
232+
safeClose(dummyState);
233+
dummyState = TestUtils.createDummyQueryFeedOperationState(
234+
ResourceType.Document,
235+
OperationType.Query,
236+
queryRequestOptions,
237+
spyClient
238+
);
239+
spyClient.queryDocuments(getCollectionLink(isNameBased), query, dummyState, Document.class).blockFirst();
240+
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
241+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
242+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
244243

245-
// Session token validation for readAll with partition query
246-
spyClient.clearCapturedRequests();
247-
queryRequestOptions = new CosmosQueryRequestOptions();
248-
safeClose(dummyState);
249-
dummyState = TestUtils.createDummyQueryFeedOperationState(
250-
ResourceType.Document,
251-
OperationType.ReadFeed,
252-
queryRequestOptions,
253-
spyClient
254-
);
255-
spyClient.readAllDocuments(
256-
getCollectionLink(isNameBased),
257-
new PartitionKey(documentCreated.getId()),
258-
dummyState,
259-
Document.class).blockFirst();
260-
assertThat(getSessionTokensInRequests().size()).isEqualTo(1);
261-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
262-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
244+
// Session token validation for readAll with partition query
245+
spyClient.clearCapturedRequests();
246+
queryRequestOptions = new CosmosQueryRequestOptions();
247+
safeClose(dummyState);
248+
dummyState = TestUtils.createDummyQueryFeedOperationState(
249+
ResourceType.Document,
250+
OperationType.ReadFeed,
251+
queryRequestOptions,
252+
spyClient
253+
);
254+
spyClient.readAllDocuments(
255+
getCollectionLink(isNameBased),
256+
new PartitionKey(documentCreated.getId()),
257+
dummyState,
258+
Document.class).blockFirst();
259+
assertThat(getSessionTokensInRequests().size()).isEqualTo(1);
260+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
261+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
263262

264-
// Session token validation for readAll with cross partition
265-
spyClient.clearCapturedRequests();
266-
queryRequestOptions = new CosmosQueryRequestOptions();
263+
// Session token validation for readAll with cross partition
264+
spyClient.clearCapturedRequests();
265+
queryRequestOptions = new CosmosQueryRequestOptions();
267266

268-
safeClose(dummyState);
269-
dummyState = TestUtils.createDummyQueryFeedOperationState(
270-
ResourceType.Document,
271-
OperationType.ReadFeed,
272-
queryRequestOptions,
273-
spyClient
274-
);
267+
safeClose(dummyState);
268+
dummyState = TestUtils.createDummyQueryFeedOperationState(
269+
ResourceType.Document,
270+
OperationType.ReadFeed,
271+
queryRequestOptions,
272+
spyClient
273+
);
275274

276-
spyClient.readDocuments(getCollectionLink(isNameBased), dummyState, Document.class).blockFirst();
277-
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
278-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
279-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
275+
spyClient.readDocuments(getCollectionLink(isNameBased), dummyState, Document.class).blockFirst();
276+
assertThat(getSessionTokensInRequests().size()).isGreaterThanOrEqualTo(1);
277+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
278+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
280279

281-
// Session token validation for readMany with cross partition
280+
// Session token validation for readMany with cross partition
281+
spyClient.clearCapturedRequests();
282+
queryRequestOptions = new CosmosQueryRequestOptions();
283+
CosmosItemIdentity cosmosItemIdentity = new CosmosItemIdentity(new PartitionKey(documentCreated.getId()), documentCreated.getId());
284+
List<CosmosItemIdentity> cosmosItemIdentities = new ArrayList<>();
285+
cosmosItemIdentities.add(cosmosItemIdentity);
286+
safeClose(dummyState);
287+
dummyState = TestUtils.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, queryRequestOptions, spyClient);
288+
spyClient.readMany(
289+
cosmosItemIdentities,
290+
getCollectionLink(isNameBased),
291+
dummyState,
292+
InternalObjectNode.class).block();
293+
assertThat(getSessionTokensInRequests().size()).isEqualTo(1);
294+
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
295+
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
296+
// session token
297+
298+
// Session token validation for create in Batch
299+
if(isNameBased) { // Batch only work with name based url
282300
spyClient.clearCapturedRequests();
283-
queryRequestOptions = new CosmosQueryRequestOptions();
284-
CosmosItemIdentity cosmosItemIdentity = new CosmosItemIdentity(new PartitionKey(documentCreated.getId()), documentCreated.getId());
285-
List<CosmosItemIdentity> cosmosItemIdentities = new ArrayList<>();
286-
cosmosItemIdentities.add(cosmosItemIdentity);
287-
safeClose(dummyState);
288-
dummyState = TestUtils
289-
.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, queryRequestOptions, spyClient);
290-
spyClient.readMany(
291-
cosmosItemIdentities,
292-
getCollectionLink(isNameBased),
293-
dummyState,
294-
InternalObjectNode.class).block();
301+
Document document = newDocument();
302+
document.set("mypk", document.getId());
303+
ItemBatchOperation<Document> itemBatchOperation = new ItemBatchOperation<Document>(CosmosItemOperationType.CREATE,
304+
documentCreated.getId(), new PartitionKey(documentCreated.getId()), new RequestOptions(), document);
305+
List<ItemBatchOperation<Document>> itemBatchOperations = new ArrayList<>();
306+
itemBatchOperations.add(itemBatchOperation);
307+
308+
Method method = SinglePartitionKeyServerBatchRequest.class.getDeclaredMethod("createBatchRequest",
309+
PartitionKey.class,
310+
List.class);
311+
method.setAccessible(true);
312+
SinglePartitionKeyServerBatchRequest serverBatchRequest =
313+
(SinglePartitionKeyServerBatchRequest) method.invoke(SinglePartitionKeyServerBatchRequest.class, new PartitionKey(document.getId()),
314+
itemBatchOperations);
315+
spyClient
316+
.executeBatchRequest(
317+
getCollectionLink(isNameBased),
318+
serverBatchRequest,
319+
new RequestOptions(),
320+
false,
321+
true)
322+
.block();
295323
assertThat(getSessionTokensInRequests().size()).isEqualTo(1);
296324
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
297325
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
298-
// session token
299-
300-
// Session token validation for create in Batch
301-
if (isNameBased) { // Batch only work with name based url
302-
spyClient.clearCapturedRequests();
303-
Document document = newDocument();
304-
document.set("mypk", document.getId());
305-
ItemBatchOperation<Document> itemBatchOperation = new ItemBatchOperation<Document>(CosmosItemOperationType.CREATE,
306-
documentCreated.getId(), new PartitionKey(documentCreated.getId()), new RequestOptions(), document);
307-
List<ItemBatchOperation<Document>> itemBatchOperations = new ArrayList<>();
308-
itemBatchOperations.add(itemBatchOperation);
309-
310-
Method method = SinglePartitionKeyServerBatchRequest.class.getDeclaredMethod("createBatchRequest",
311-
PartitionKey.class,
312-
List.class);
313-
method.setAccessible(true);
314-
SinglePartitionKeyServerBatchRequest serverBatchRequest =
315-
(SinglePartitionKeyServerBatchRequest) method.invoke(SinglePartitionKeyServerBatchRequest.class, new PartitionKey(document.getId()),
316-
itemBatchOperations);
317-
spyClient
318-
.executeBatchRequest(
319-
getCollectionLink(isNameBased),
320-
serverBatchRequest,
321-
new RequestOptions(),
322-
false,
323-
true)
324-
.block();
325-
assertThat(getSessionTokensInRequests().size()).isEqualTo(1);
326-
assertThat(getSessionTokensInRequests().get(0)).isNotEmpty();
327-
assertThat(getSessionTokensInRequests().get(0)).doesNotContain(","); // making sure we have only one scope session token
328-
}
329-
} finally {
330-
safeClose(dummyState);
331326
}
327+
328+
safeClose(dummyState);
332329
}
333330

334331
@Test(groups = { "fast" }, timeOut = TIMEOUT, dataProvider = "sessionTestArgProvider")

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -229,9 +229,8 @@ public void crossPartitionQuery() {
229229
options.setMaxDegreeOfParallelism(-1);
230230
ModelBridgeInternal.setQueryRequestOptionsMaxItemCount(options, 100);
231231

232-
QueryFeedOperationState dummyState = TestUtils
233-
.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
234-
232+
QueryFeedOperationState dummyState =
233+
TestUtils.createDummyQueryFeedOperationState(ResourceType.Document, OperationType.Query, options, client);
235234
try {
236235
Flux<FeedResponse<Document>> results = client.queryDocuments(
237236
getCollectionLink(),
@@ -245,6 +244,7 @@ public void crossPartitionQuery() {
245244

246245
validateQuerySuccess(results, validator, QUERY_TIMEOUT);
247246
validateNoDocumentQueryOperationThroughGateway();
247+
248248
// validates only the first query for fetching query plan goes to gateway.
249249
assertThat(client.getCapturedRequests().stream().filter(r -> r.getResourceType() == ResourceType.Document)).hasSize(1);
250250
} finally {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,15 @@ public static void setTransportClient(StoreReader storeReader, TransportClient t
317317
set(storeReader, transportClient, "transportClient");
318318
}
319319

320+
public static void setTransportClient(CosmosClient client, TransportClient transportClient) {
321+
StoreClient storeClient = getStoreClient((RxDocumentClientImpl) CosmosBridgeInternal.getAsyncDocumentClient(client));
322+
set(storeClient, transportClient, "transportClient");
323+
ReplicatedResourceClient replicatedResClient = getReplicatedResourceClient(storeClient);
324+
ConsistencyWriter writer = getConsistencyWriter(replicatedResClient);
325+
set(replicatedResClient, transportClient, "transportClient");
326+
set(writer, transportClient, "transportClient");
327+
}
328+
320329
public static TransportClient getTransportClient(ReplicatedResourceClient replicatedResourceClient) {
321330
return get(TransportClient.class, replicatedResourceClient, "transportClient");
322331
}

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1138,6 +1138,11 @@ public URI serviceEndpoint() {
11381138
return null;
11391139
}
11401140

1141+
@Override
1142+
public URI serverKeyUsedAsActualRemoteAddress() {
1143+
return this.remoteURI;
1144+
}
1145+
11411146
@Override
11421147
public void injectConnectionErrors(String ruleId, double threshold, Class<?> eventType) {
11431148
throw new NotImplementedException("injectConnectionErrors is not supported in FakeEndpoint");

0 commit comments

Comments
 (0)