public async Task TestMoveNextWithEmptyPagesAndSplitAsync(string initialContinuationToken)
{
int maxPageSize = 5;
List <MockPartitionResponse[]> mockResponsesScenario = MockQueryFactory.GetSplitScenarios();
foreach (MockPartitionResponse[] mockResponse in mockResponsesScenario)
{
Mock <CosmosQueryClient> mockQueryClient = new Mock <CosmosQueryClient>();
IList <ToDoItem> allItems = MockQueryFactory.GenerateAndMockResponse(
mockQueryClient,
isOrderByQuery: false,
sqlQuerySpec: MockQueryFactory.DefaultQuerySpec,
containerRid: MockQueryFactory.DefaultCollectionRid,
initContinuationToken: initialContinuationToken,
maxPageSize: maxPageSize,
mockResponseForSinglePartition: mockResponse,
cancellationTokenForMocks: this.cancellationToken);
CosmosQueryContext context = MockQueryFactory.CreateContext(
mockQueryClient.Object);
ItemProducerTree itemProducerTree = new ItemProducerTree(
context,
MockQueryFactory.DefaultQuerySpec,
mockResponse[0].PartitionKeyRange,
MockItemProducerFactory.DefaultTreeProduceAsyncCompleteDelegate,
new ParallelItemProducerTreeComparer(),
CosmosElementEqualityComparer.Value,
true,
MockQueryFactory.DefaultCollectionRid,
maxPageSize,
initialContinuationToken: initialContinuationToken);
Assert.IsTrue(itemProducerTree.HasMoreResults);
List <ToDoItem> itemsRead = new List <ToDoItem>();
while ((await itemProducerTree.MoveNextAsync(this.cancellationToken)).successfullyMovedNext)
{
Assert.IsTrue(itemProducerTree.HasMoreResults);
if (itemProducerTree.Current != null)
{
string jsonValue = itemProducerTree.Current.ToString();
ToDoItem item = JsonConvert.DeserializeObject <ToDoItem>(jsonValue);
itemsRead.Add(item);
}
}
Assert.IsFalse(itemProducerTree.HasMoreResults);
Assert.AreEqual(allItems.Count, itemsRead.Count);
List <ToDoItem> exepected = allItems.OrderBy(x => x.id).ToList();
List <ToDoItem> actual = itemsRead.OrderBy(x => x.id).ToList();
CollectionAssert.AreEqual(exepected, actual, new ToDoItemComparer());
}
}
public async Task TestMoveNextWithEmptyPagesAsync(string initialContinuationToken)
{
int maxPageSize = 5;
List <MockPartitionResponse[]> mockResponses = MockQueryFactory.GetAllCombinationWithEmptyPage();
foreach (MockPartitionResponse[] mockResponse in mockResponses)
{
Mock <CosmosQueryClient> mockQueryClient = new Mock <CosmosQueryClient>();
IList <ToDoItem> allItems = MockQueryFactory.GenerateAndMockResponse(
mockQueryClient,
isOrderByQuery: false,
sqlQuerySpec: MockQueryFactory.DefaultQuerySpec,
containerRid: MockQueryFactory.DefaultCollectionRid,
initContinuationToken: initialContinuationToken,
maxPageSize: maxPageSize,
mockResponseForSinglePartition: mockResponse,
cancellationTokenForMocks: this.cancellationToken);
CosmosQueryContext context = MockQueryFactory.CreateContext(
mockQueryClient.Object);
ItemProducerTree itemProducerTree = new ItemProducerTree(
queryContext: context,
querySpecForInit: MockQueryFactory.DefaultQuerySpec,
partitionKeyRange: mockResponse[0].PartitionKeyRange,
produceAsyncCompleteCallback: MockItemProducerFactory.DefaultTreeProduceAsyncCompleteDelegate,
itemProducerTreeComparer: new ParallelItemProducerTreeComparer(),
equalityComparer: CosmosElementEqualityComparer.Value,
deferFirstPage: true,
collectionRid: MockQueryFactory.DefaultCollectionRid,
initialPageSize: maxPageSize,
initialContinuationToken: initialContinuationToken);
Assert.IsTrue(itemProducerTree.HasMoreResults);
List <ToDoItem> itemsRead = new List <ToDoItem>();
while ((await itemProducerTree.MoveNextAsync(this.cancellationToken)).successfullyMovedNext)
{
Assert.IsTrue(itemProducerTree.HasMoreResults);
string jsonValue = itemProducerTree.Current.ToString();
ToDoItem item = JsonConvert.DeserializeObject <ToDoItem>(jsonValue);
itemsRead.Add(item);
}
Assert.IsFalse(itemProducerTree.HasMoreResults);
Assert.AreEqual(allItems.Count, itemsRead.Count);
CollectionAssert.AreEqual(itemsRead, allItems.ToList(), new ToDoItemComparer());
}
}
public async Task TestItemProducerTreeWithFailure()
{
int callBackCount = 0;
Mock <CosmosQueryContext> mockQueryContext = new Mock <CosmosQueryContext>();
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("Select * from t");
PartitionKeyRange partitionKeyRange = new PartitionKeyRange {
Id = "0", MinInclusive = "A", MaxExclusive = "B"
};
Action <ItemProducerTree, int, double, QueryMetrics, long, CancellationToken> produceAsyncCompleteCallback = (
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
QueryMetrics queryMetrics,
long responseLengthBytes,
CancellationToken token) => { callBackCount++; };
Mock <IComparer <ItemProducerTree> > comparer = new Mock <IComparer <ItemProducerTree> >();
Mock <IEqualityComparer <CosmosElement> > cosmosElementComparer = new Mock <IEqualityComparer <CosmosElement> >();
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
IEnumerable <CosmosElement> cosmosElements = new List <CosmosElement>()
{
new Mock <CosmosElement>(CosmosElementType.Object).Object
};
CosmosQueryResponseMessageHeaders headers = new CosmosQueryResponseMessageHeaders("TestToken", null)
{
ActivityId = "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE",
RequestCharge = 42
};
mockQueryContext.Setup(x => x.ExecuteQueryAsync(sqlQuerySpec, cancellationTokenSource.Token, It.IsAny <Action <CosmosRequestMessage> >())).Returns(
Task.FromResult(QueryResponse.CreateSuccess(cosmosElements, 1, 500, headers)));
ItemProducerTree itemProducerTree = new ItemProducerTree(
queryContext: mockQueryContext.Object,
querySpecForInit: sqlQuerySpec,
partitionKeyRange: partitionKeyRange,
produceAsyncCompleteCallback: produceAsyncCompleteCallback,
itemProducerTreeComparer: comparer.Object,
equalityComparer: cosmosElementComparer.Object,
deferFirstPage: false,
collectionRid: "collectionRid",
initialContinuationToken: null,
initialPageSize: 50);
// Buffer to success responses
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// Buffer a failure
mockQueryContext.Setup(x => x.ExecuteQueryAsync(sqlQuerySpec, cancellationTokenSource.Token, It.IsAny <Action <CosmosRequestMessage> >())).Returns(
Task.FromResult(QueryResponse.CreateFailure(headers, HttpStatusCode.InternalServerError, null, "Error message", null)));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// First item should be a success
var result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Second item should be a success
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Third item should be a failure
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
// Try to buffer after failure. It should return the previous cached failure and not try to buffer again.
mockQueryContext.Setup(x => x.ExecuteQueryAsync(sqlQuerySpec, cancellationTokenSource.Token, It.IsAny <Action <CosmosRequestMessage> >())).
Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again"));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
}
public async Task TestItemProducerTreeWithFailure()
{
int callBackCount = 0;
Mock <CosmosQueryContext> mockQueryContext = new Mock <CosmosQueryContext>();
SqlQuerySpec sqlQuerySpec = new SqlQuerySpec("Select * from t");
PartitionKeyRange partitionKeyRange = new PartitionKeyRange {
Id = "0", MinInclusive = "A", MaxExclusive = "B"
};
ItemProducerTree.ProduceAsyncCompleteDelegate produceAsyncCompleteCallback = (
ItemProducerTree producer,
int itemsBuffered,
double resourceUnitUsage,
IReadOnlyCollection <QueryPageDiagnostics> queryPageDiagnostics,
long responseLengthBytes,
CancellationToken token) =>
{ callBackCount++; };
Mock <IComparer <ItemProducerTree> > comparer = new Mock <IComparer <ItemProducerTree> >();
Mock <IEqualityComparer <CosmosElement> > cosmosElementComparer = new Mock <IEqualityComparer <CosmosElement> >();
CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
IReadOnlyList <CosmosElement> cosmosElements = new List <CosmosElement>()
{
new Mock <CosmosElement>(CosmosElementType.Object).Object
};
QueryPageDiagnostics diagnostics = new QueryPageDiagnostics(
partitionKeyRangeId: "0",
queryMetricText: "SomeRandomQueryMetricText",
indexUtilizationText: null,
requestDiagnostics: new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.OK,
subStatusCode: SubStatusCodes.Unknown,
requestCharge: 42,
errorMessage: null,
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
clientSideRequestStatistics: null),
schedulingStopwatch: new SchedulingStopwatch());
IReadOnlyCollection <QueryPageDiagnostics> pageDiagnostics = new List <QueryPageDiagnostics>()
{
diagnostics
};
mockQueryContext.Setup(x => x.ContainerResourceId).Returns("MockCollectionRid");
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateSuccess(
result: cosmosElements,
requestCharge: 42,
activityId: "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE",
diagnostics: pageDiagnostics,
responseLengthBytes: 500,
disallowContinuationTokenMessage: null,
continuationToken: "TestToken")));
ItemProducerTree itemProducerTree = new ItemProducerTree(
queryContext: mockQueryContext.Object,
querySpecForInit: sqlQuerySpec,
partitionKeyRange: partitionKeyRange,
produceAsyncCompleteCallback: produceAsyncCompleteCallback,
itemProducerTreeComparer: comparer.Object,
equalityComparer: cosmosElementComparer.Object,
deferFirstPage: false,
collectionRid: "collectionRid",
initialContinuationToken: null,
initialPageSize: 50);
// Buffer to success responses
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
diagnostics = new QueryPageDiagnostics(
partitionKeyRangeId: "0",
queryMetricText: null,
indexUtilizationText: null,
requestDiagnostics: new PointOperationStatistics(
Guid.NewGuid().ToString(),
System.Net.HttpStatusCode.InternalServerError,
subStatusCode: SubStatusCodes.Unknown,
requestCharge: 10.2,
errorMessage: "Error message",
method: HttpMethod.Post,
requestUri: new Uri("http://localhost.com"),
clientSideRequestStatistics: null),
schedulingStopwatch: new SchedulingStopwatch());
pageDiagnostics = new List <QueryPageDiagnostics>()
{
diagnostics
};
// Buffer a failure
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).Returns(
Task.FromResult(QueryResponseCore.CreateFailure(
statusCode: HttpStatusCode.InternalServerError,
subStatusCodes: null,
errorMessage: "Error message",
requestCharge: 10.2,
activityId: Guid.NewGuid().ToString(),
diagnostics: pageDiagnostics)));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
// First item should be a success
(bool successfullyMovedNext, QueryResponseCore? failureResponse)result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Second item should be a success
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsTrue(result.successfullyMovedNext);
Assert.IsNull(result.failureResponse);
Assert.IsTrue(itemProducerTree.HasMoreResults);
// Third item should be a failure
result = await itemProducerTree.MoveNextAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
// Try to buffer after failure. It should return the previous cached failure and not try to buffer again.
mockQueryContext.Setup(x => x.ExecuteQueryAsync(
sqlQuerySpec,
It.IsAny <string>(),
It.IsAny <PartitionKeyRangeIdentity>(),
It.IsAny <bool>(),
It.IsAny <int>(),
It.IsAny <SchedulingStopwatch>(),
cancellationTokenSource.Token)).
Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again"));
await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);
Assert.IsFalse(result.successfullyMovedNext);
Assert.IsNotNull(result.failureResponse);
Assert.IsFalse(itemProducerTree.HasMoreResults);
}
