Beispiel #1
0
        public void ValidateDiagnosticsAppendContext()
        {
            CosmosDiagnosticsContext cosmosDiagnostics = new CosmosDiagnosticsContextCore(
                nameof(ValidateDiagnosticsAppendContext),
                "MyCustomUserAgentString");
            CosmosDiagnosticsContext cosmosDiagnostics2;

            using (cosmosDiagnostics.GetOverallScope())
            {
                // Test all the different operations on diagnostics context
                using (cosmosDiagnostics.CreateScope("ValidateScope"))
                {
                    Thread.Sleep(TimeSpan.FromSeconds(2));
                }

                cosmosDiagnostics2 = new CosmosDiagnosticsContextCore(
                    nameof(ValidateDiagnosticsAppendContext),
                    "MyCustomUserAgentString");
                cosmosDiagnostics2.GetOverallScope().Dispose();

                using (cosmosDiagnostics.CreateScope("CosmosDiagnostics2Scope"))
                {
                    Thread.Sleep(TimeSpan.FromMilliseconds(100));
                }

                cosmosDiagnostics2.AddDiagnosticsInternal(cosmosDiagnostics);
            }

            string diagnostics = cosmosDiagnostics2.ToString();

            Assert.IsTrue(diagnostics.Contains("MyCustomUserAgentString"));
            Assert.IsTrue(diagnostics.Contains("ValidateScope"));
            Assert.IsTrue(diagnostics.Contains("CosmosDiagnostics2Scope"));
        }
Beispiel #2
0
        private (Func <CosmosElement, Task <TryCatch <IDocumentQueryExecutionComponent> > >, QueryResponseCore) SetupBaseContextToVerifyFailureScenario()
        {
            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();

            diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
                                                          Guid.NewGuid().ToString(),
                                                          System.Net.HttpStatusCode.Unauthorized,
                                                          subStatusCode: SubStatusCodes.PartitionKeyMismatch,
                                                          responseTimeUtc: DateTime.UtcNow,
                                                          requestCharge: 4,
                                                          errorMessage: null,
                                                          method: HttpMethod.Post,
                                                          requestUri: new Uri("http://localhost.com"),
                                                          requestSessionToken: null,
                                                          responseSessionToken: null));
            IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
            {
                new QueryPageDiagnostics(
                    "0",
                    "SomeQueryMetricText",
                    "SomeIndexUtilText",
                    diagnosticsContext,
                    new SchedulingStopwatch())
            };

            QueryResponseCore failure = QueryResponseCore.CreateFailure(
                System.Net.HttpStatusCode.Unauthorized,
                SubStatusCodes.PartitionKeyMismatch,
                new CosmosException(
                    statusCodes: HttpStatusCode.Unauthorized,
                    message: "Random error message",
                    subStatusCode: default,
        public static QueryResponseCore CreateQueryResponse(
            IList <ToDoItem> items,
            bool isOrderByQuery,
            string continuationToken,
            string collectionRid)
        {
            MemoryStream memoryStream;
            string       json;

            if (isOrderByQuery)
            {
                memoryStream = SerializeForOrderByQuery(items);
                using (StreamReader sr = new StreamReader(SerializeForOrderByQuery(items)))
                {
                    json = sr.ReadToEnd();
                }
            }
            else
            {
                memoryStream = (MemoryStream)MockCosmosUtil.Serializer.ToStream <IList <ToDoItem> >(items);
            }

            long responseLengthBytes = memoryStream.Length;

            IJsonNavigator           jsonNavigator      = JsonNavigator.Create(memoryStream.ToArray());
            IJsonNavigatorNode       jsonNavigatorNode  = jsonNavigator.GetRootNode();
            CosmosArray              cosmosArray        = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);
            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();

            diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
                                                          activityId: Guid.NewGuid().ToString(),
                                                          statusCode: HttpStatusCode.OK,
                                                          subStatusCode: SubStatusCodes.Unknown,
                                                          responseTimeUtc: DateTime.UtcNow,
                                                          requestCharge: 4,
                                                          errorMessage: null,
                                                          method: HttpMethod.Post,
                                                          requestUri: "http://localhost.com",
                                                          requestSessionToken: null,
                                                          responseSessionToken: null));
            IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
            {
                new QueryPageDiagnostics(
                    Guid.NewGuid(),
                    "0",
                    "SomeQueryMetricText",
                    "SomeIndexUtilText",
                    diagnosticsContext)
            };

            QueryResponseCore message = QueryResponseCore.CreateSuccess(
                result: cosmosArray,
                requestCharge: 4,
                activityId: Guid.NewGuid().ToString(),
                responseLengthBytes: responseLengthBytes,
                disallowContinuationTokenMessage: null,
                continuationToken: continuationToken);

            return(message);
        }
Beispiel #4
0
        public void ValidateDiagnosticsAppendContext()
        {
            CosmosDiagnosticsContext cosmosDiagnostics = new CosmosDiagnosticsContextCore(
                nameof(ValidateDiagnosticsAppendContext),
                "MyCustomUserAgentString");
            CosmosDiagnosticsContext cosmosDiagnostics2;

            using (cosmosDiagnostics.GetOverallScope())
            {
                bool insertIntoDiagnostics1 = true;
                bool isInsertDiagnostics    = false;
                // Start a background thread and ensure that no exception occurs even if items are getting added to the context
                // when 2 contexts are appended.
                Task.Run(() =>
                {
                    isInsertDiagnostics = true;
                    CosmosSystemInfo cosmosSystemInfo = new CosmosSystemInfo(
                        cpuLoadHistory: new Documents.Rntbd.CpuLoadHistory(new List <Documents.Rntbd.CpuLoad>().AsReadOnly(), TimeSpan.FromSeconds(1)));
                    while (insertIntoDiagnostics1)
                    {
                        Task.Delay(TimeSpan.FromMilliseconds(1)).Wait();
                        cosmosDiagnostics.AddDiagnosticsInternal(cosmosSystemInfo);
                    }
                });

                while (!isInsertDiagnostics)
                {
                    Task.Delay(TimeSpan.FromMilliseconds(5)).Wait();
                }

                // Test all the different operations on diagnostics context
                using (cosmosDiagnostics.CreateScope("ValidateScope"))
                {
                    Thread.Sleep(TimeSpan.FromMilliseconds(3));
                }

                cosmosDiagnostics2 = new CosmosDiagnosticsContextCore(
                    nameof(ValidateDiagnosticsAppendContext),
                    "MyCustomUserAgentString");
                cosmosDiagnostics2.GetOverallScope().Dispose();

                using (cosmosDiagnostics.CreateScope("CosmosDiagnostics2Scope"))
                {
                    Thread.Sleep(TimeSpan.FromMilliseconds(3));
                }

                cosmosDiagnostics2.AddDiagnosticsInternal(cosmosDiagnostics);

                // Stop the background inserts
                insertIntoDiagnostics1 = false;
            }

            string diagnostics = cosmosDiagnostics2.ToString();

            Assert.IsTrue(diagnostics.Contains("MyCustomUserAgentString"));
            Assert.IsTrue(diagnostics.Contains("ValidateScope"));
            Assert.IsTrue(diagnostics.Contains("CosmosDiagnostics2Scope"));
        }
        public static (QueryResponseCore queryResponse, IList <ToDoItem> items) Create(
            string itemIdPrefix,
            string continuationToken,
            string collectionRid,
            int itemCount = 50)
        {
            // Use -1 to represent a split response
            if (itemCount == QueryResponseMessageFactory.SPLIT)
            {
                return(CreateSplitResponse(collectionRid), new List <ToDoItem>().AsReadOnly());
            }

            IList <ToDoItem> items               = ToDoItem.CreateItems(itemCount, itemIdPrefix);
            MemoryStream     memoryStream        = (MemoryStream)MockCosmosUtil.Serializer.ToStream <IList <ToDoItem> >(items);
            long             responseLengthBytes = memoryStream.Length;

            IJsonNavigator     jsonNavigator     = JsonNavigator.Create(memoryStream.ToArray());
            IJsonNavigatorNode jsonNavigatorNode = jsonNavigator.GetRootNode();
            CosmosArray        cosmosArray       = CosmosArray.Create(jsonNavigator, jsonNavigatorNode);

            double requestCharge = 42;
            string activityId    = Guid.NewGuid().ToString();
            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();

            diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
                                                          activityId: Guid.NewGuid().ToString(),
                                                          statusCode: HttpStatusCode.OK,
                                                          subStatusCode: SubStatusCodes.Unknown,
                                                          responseTimeUtc: DateTime.UtcNow,
                                                          requestCharge: requestCharge,
                                                          errorMessage: null,
                                                          method: HttpMethod.Post,
                                                          requestUri: "http://localhost.com",
                                                          requestSessionToken: null,
                                                          responseSessionToken: null));
            IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
            {
                new QueryPageDiagnostics(
                    Guid.NewGuid(),
                    "0",
                    "SomeQueryMetricText",
                    "SomeIndexUtilText",
                    diagnosticsContext)
            };

            QueryResponseCore message = QueryResponseCore.CreateSuccess(
                result: cosmosArray,
                continuationToken: continuationToken,
                disallowContinuationTokenMessage: null,
                activityId: activityId,
                requestCharge: requestCharge,
                responseLengthBytes: responseLengthBytes);

            return(message, items);
        }
        public void TestCosmosDiagnosticScope()
        {
            (BackendMetricsExtractor.ParseFailureReason parseFailureReason, BackendMetrics extractedBackendMetrics)notFoundResult = MockCosmosDiagnosticScope.Accept(BackendMetricsExtractor.Singleton);
            Assert.AreEqual(BackendMetricsExtractor.ParseFailureReason.MetricsNotFound, notFoundResult.parseFailureReason);

            CosmosDiagnosticsContext contextWithQueryMetrics = new CosmosDiagnosticsContextCore();

            contextWithQueryMetrics.AddDiagnosticsInternal(MockQueryPageDiagnostics);
            (BackendMetricsExtractor.ParseFailureReason parseFailureReason, BackendMetrics extractedBackendMetrics)foundResult = contextWithQueryMetrics.Accept(BackendMetricsExtractor.Singleton);
            Assert.AreEqual(BackendMetricsExtractor.ParseFailureReason.None, foundResult.parseFailureReason);
            Assert.AreEqual(BackendMetricsTests.MockBackendMetrics.IndexLookupTime, foundResult.extractedBackendMetrics.IndexLookupTime);
        }
        public static QueryResponseCore CreateFailureResponse(
            HttpStatusCode httpStatusCode,
            SubStatusCodes subStatusCodes,
            string errorMessage)
        {
            string acitivityId = Guid.NewGuid().ToString();
            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();

            diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
                                                          acitivityId,
                                                          System.Net.HttpStatusCode.Gone,
                                                          subStatusCode: SubStatusCodes.PartitionKeyRangeGone,
                                                          responseTimeUtc: DateTime.UtcNow,
                                                          requestCharge: 10.4,
                                                          errorMessage: null,
                                                          method: HttpMethod.Post,
                                                          requestUri: "http://localhost.com",
                                                          requestSessionToken: null,
                                                          responseSessionToken: null));
            IReadOnlyCollection <QueryPageDiagnostics> diagnostics = new List <QueryPageDiagnostics>()
            {
                new QueryPageDiagnostics(
                    Guid.NewGuid(),
                    "0",
                    "SomeQueryMetricText",
                    "SomeIndexUtilText",
                    diagnosticsContext)
            };

            QueryResponseCore splitResponse = QueryResponseCore.CreateFailure(
                statusCode: httpStatusCode,
                subStatusCodes: subStatusCodes,
                cosmosException: CosmosExceptionFactory.Create(
                    statusCode: httpStatusCode,
                    subStatusCode: (int)subStatusCodes,
                    message: errorMessage,
                    stackTrace: new System.Diagnostics.StackTrace().ToString(),
                    activityId: acitivityId,
                    requestCharge: 10.4,
                    retryAfter: default,
        public void ToResponseMessage_MapsProperties()
        {
            PointOperationStatistics pointOperationStatistics = new PointOperationStatistics(
                activityId: Guid.NewGuid().ToString(),
                statusCode: HttpStatusCode.OK,
                subStatusCode: SubStatusCodes.Unknown,
                responseTimeUtc: DateTime.UtcNow,
                requestCharge: 0,
                errorMessage: string.Empty,
                method: HttpMethod.Get,
                requestUri: "http://localhost",
                requestSessionToken: null,
                responseSessionToken: null);
            CosmosDiagnosticsContext context = new CosmosDiagnosticsContextCore();

            context.AddDiagnosticsInternal(pointOperationStatistics);

            TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(HttpStatusCode.OK)
            {
                ResourceStream     = new MemoryStream(new byte[] { 0x41, 0x42 }, index: 0, count: 2, writable: false, publiclyVisible: true),
                ETag               = "1234",
                SubStatusCode      = SubStatusCodes.CompletingSplit,
                RetryAfter         = TimeSpan.FromSeconds(10),
                RequestCharge      = 4.3,
                DiagnosticsContext = context
            };

            ResponseMessage response = result.ToResponseMessage();

            Assert.AreEqual(result.ResourceStream, response.Content);
            Assert.AreEqual(result.SubStatusCode, response.Headers.SubStatusCode);
            Assert.AreEqual(result.RetryAfter, response.Headers.RetryAfter);
            Assert.AreEqual(result.StatusCode, response.StatusCode);
            Assert.AreEqual(result.RequestCharge, response.Headers.RequestCharge);
            string diagnostics = response.Diagnostics.ToString();

            Assert.IsNotNull(diagnostics);
            Assert.IsTrue(diagnostics.Contains(pointOperationStatistics.ActivityId));
        }
Beispiel #9
0
        public void ValidateRetriableRequestsCount()
        {
            CosmosDiagnosticsContext cosmosDiagnostics = new CosmosDiagnosticsContextCore(
                nameof(ValidateRetriableRequestsCount),
                "cosmos-netstandard-sdk");

            cosmosDiagnostics.GetOverallScope().Dispose();
            using (cosmosDiagnostics.GetOverallScope())
            {
                // Test all the different operations on diagnostics context
                using (cosmosDiagnostics.CreateScope("ValidateScope"))
                {
                    Assert.AreEqual(0, cosmosDiagnostics.GetRetriableResponseCount());
                    Assert.AreEqual(0, cosmosDiagnostics.GetFailedResponseCount());
                    Assert.AreEqual(0, cosmosDiagnostics.GetTotalResponseCount());

                    cosmosDiagnostics.AddDiagnosticsInternal(new PointOperationStatistics(
                                                                 new Guid("692ab2f2-41ba-486b-aad7-8c7c6c52379f").ToString(),
                                                                 (HttpStatusCode)429,
                                                                 Documents.SubStatusCodes.Unknown,
                                                                 DateTime.UtcNow,
                                                                 42,
                                                                 null,
                                                                 HttpMethod.Get,
                                                                 "http://MockUri.com",
                                                                 null,
                                                                 null));

                    Assert.AreEqual(1, cosmosDiagnostics.GetRetriableResponseCount());
                    Assert.AreEqual(1, cosmosDiagnostics.GetFailedResponseCount());
                    Assert.AreEqual(1, cosmosDiagnostics.GetTotalResponseCount());

                    cosmosDiagnostics.AddDiagnosticsInternal(
                        new StoreResponseStatistics(
                            DateTime.UtcNow,
                            DateTime.UtcNow,
                            StoreResult.CreateStoreResult(
                                new StoreResponse {
                        Status = 449
                    }, null, false, false),
                            ResourceType.Document,
                            OperationType.Delete,
                            new Uri("http://MockUri.com")));

                    Assert.AreEqual(2, cosmosDiagnostics.GetRetriableResponseCount());
                    Assert.AreEqual(2, cosmosDiagnostics.GetFailedResponseCount());
                    Assert.AreEqual(2, cosmosDiagnostics.GetTotalResponseCount());

                    cosmosDiagnostics.AddDiagnosticsInternal(
                        new StoreResponseStatistics(
                            DateTime.UtcNow,
                            DateTime.UtcNow,
                            StoreResult.CreateStoreResult(
                                new StoreResponse {
                        Status = 503
                    }, null, false, false),
                            ResourceType.Document,
                            OperationType.Delete,
                            new Uri("http://MockUri.com")));

                    Assert.AreEqual(2, cosmosDiagnostics.GetRetriableResponseCount());
                    Assert.AreEqual(3, cosmosDiagnostics.GetFailedResponseCount());
                    Assert.AreEqual(3, cosmosDiagnostics.GetTotalResponseCount());
                }

                Assert.AreEqual(2, cosmosDiagnostics.GetRetriableResponseCount());
                Assert.AreEqual(3, cosmosDiagnostics.GetFailedResponseCount());
                Assert.AreEqual(3, cosmosDiagnostics.GetTotalResponseCount());
            }

            Assert.AreEqual(2, cosmosDiagnostics.GetRetriableResponseCount());
            Assert.AreEqual(3, cosmosDiagnostics.GetFailedResponseCount());
            Assert.AreEqual(3, cosmosDiagnostics.GetTotalResponseCount());

            CosmosDiagnosticsContext cosmosDiagnostics2 = new CosmosDiagnosticsContextCore(
                nameof(ValidateRetriableRequestsCount),
                "cosmos-netstandard-sdk");

            cosmosDiagnostics2.AddDiagnosticsInternal(cosmosDiagnostics);

            Assert.AreEqual(2, cosmosDiagnostics2.GetRetriableResponseCount());
            Assert.AreEqual(3, cosmosDiagnostics2.GetFailedResponseCount());
            Assert.AreEqual(3, cosmosDiagnostics2.GetTotalResponseCount());
        }
Beispiel #10
0
        public void ValidateDiagnosticsContext()
        {
            CosmosDiagnosticsContext cosmosDiagnostics = new CosmosDiagnosticsContextCore(
                nameof(ValidateDiagnosticsContext),
                "cosmos-netstandard-sdk");

            cosmosDiagnostics.GetOverallScope().Dispose();
            string diagnostics = cosmosDiagnostics.ToString();

            //Test the default user agent string
            JObject jObject = JObject.Parse(diagnostics);
            JToken  summary = jObject["Summary"];

            Assert.IsTrue(summary["UserAgent"].ToString().Contains("cosmos-netstandard-sdk"), "Diagnostics should have user agent string");

            cosmosDiagnostics = new CosmosDiagnosticsContextCore(
                nameof(ValidateDiagnosticsContext),
                "MyCustomUserAgentString");
            using (cosmosDiagnostics.GetOverallScope())
            {
                // Test all the different operations on diagnostics context
                Thread.Sleep(TimeSpan.FromSeconds(1));
                using (cosmosDiagnostics.CreateScope("ValidateScope"))
                {
                    Thread.Sleep(TimeSpan.FromSeconds(1));
                    cosmosDiagnostics.AddDiagnosticsInternal(new PointOperationStatistics(
                                                                 new Guid("692ab2f2-41ba-486b-aad7-8c7c6c52379f").ToString(),
                                                                 (HttpStatusCode)429,
                                                                 Documents.SubStatusCodes.Unknown,
                                                                 DateTime.UtcNow,
                                                                 42,
                                                                 null,
                                                                 HttpMethod.Get,
                                                                 "http://MockUri.com",
                                                                 null,
                                                                 null));
                }

                using (cosmosDiagnostics.CreateScope("SuccessScope"))
                {
                    cosmosDiagnostics.AddDiagnosticsInternal(new PointOperationStatistics(
                                                                 new Guid("de09baab-71a4-4897-a163-470711c93ed3").ToString(),
                                                                 HttpStatusCode.OK,
                                                                 Documents.SubStatusCodes.Unknown,
                                                                 DateTime.UtcNow,
                                                                 42,
                                                                 null,
                                                                 HttpMethod.Get,
                                                                 "http://MockUri.com",
                                                                 null,
                                                                 null));
                }
            }

            string result = cosmosDiagnostics.ToString();

            string regex = @"\{""DiagnosticVersion"":""2"",""Summary"":\{""StartUtc"":"".+Z"",""TotalElapsedTimeInMs"":.+,""UserAgent"":""MyCustomUserAgentString"",""TotalRequestCount"":2,""FailedRequestCount"":1\},""Context"":\[\{""Id"":""ValidateScope"",""ElapsedTimeInMs"":.+\},\{""Id"":""PointOperationStatistics"",""ActivityId"":""692ab2f2-41ba-486b-aad7-8c7c6c52379f"",""ResponseTimeUtc"":"".+Z"",""StatusCode"":429,""SubStatusCode"":0,""RequestCharge"":42.0,""RequestUri"":""http://MockUri.com"",""RequestSessionToken"":null,""ResponseSessionToken"":null\},\{""Id"":""SuccessScope"",""ElapsedTimeInMs"":.+\},\{""Id"":""PointOperationStatistics"",""ActivityId"":""de09baab-71a4-4897-a163-470711c93ed3"",""ResponseTimeUtc"":"".+Z"",""StatusCode"":200,""SubStatusCode"":0,""RequestCharge"":42.0,""RequestUri"":""http://MockUri.com"",""RequestSessionToken"":null,""ResponseSessionToken"":null\}\]\}";

            Assert.IsTrue(Regex.IsMatch(result, regex), $"regex: {regex} result: {result}");

            JToken jToken = JToken.Parse(result);
            double total  = jToken["Summary"]["TotalElapsedTimeInMs"].ToObject <double>();

            Assert.IsTrue(total > TimeSpan.FromSeconds(2).TotalMilliseconds);
            double overalScope = jToken["Context"][0]["ElapsedTimeInMs"].ToObject <double>();

            Assert.IsTrue(overalScope < total);
            Assert.IsTrue(overalScope > TimeSpan.FromSeconds(1).TotalMilliseconds);
            double innerScope = jToken["Context"][2]["ElapsedTimeInMs"].ToObject <double>();

            Assert.IsTrue(innerScope > 0);
        }
Beispiel #11
0
        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"
            };

            void produceAsyncCompleteCallback(
                ItemProducerTree producer,
                int itemsBuffered,
                double resourceUnitUsage,
                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
            };

            CosmosDiagnosticsContext diagnosticsContext = new CosmosDiagnosticsContextCore();

            diagnosticsContext.AddDiagnosticsInternal(new PointOperationStatistics(
                                                          Guid.NewGuid().ToString(),
                                                          System.Net.HttpStatusCode.OK,
                                                          subStatusCode: SubStatusCodes.Unknown,
                                                          responseTimeUtc: DateTime.UtcNow,
                                                          requestCharge: 42,
                                                          errorMessage: null,
                                                          method: HttpMethod.Post,
                                                          requestUri: new Uri("http://localhost.com"),
                                                          requestSessionToken: null,
                                                          responseSessionToken: null));

            QueryPageDiagnostics diagnostics = new QueryPageDiagnostics(
                clientQueryCorrelationId: Guid.NewGuid(),
                partitionKeyRangeId: "0",
                queryMetricText: "SomeRandomQueryMetricText",
                indexUtilizationText: null,
                diagnosticsContext: diagnosticsContext);
            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>(),
                                       cancellationTokenSource.Token)).Returns(
                Task.FromResult(QueryResponseCore.CreateSuccess(
                                    result: cosmosElements,
                                    requestCharge: 42,
                                    activityId: "AA470D71-6DEF-4D61-9A08-272D8C9ABCFE",
                                    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,
                testSettings: new TestInjections(simulate429s: false, simulateEmptyPages: false),
                deferFirstPage: false,
                collectionRid: "collectionRid",
                initialContinuationToken: null,
                initialPageSize: 50);

            // Buffer to success responses
            await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);

            await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);

            CosmosDiagnosticsContext diagnosticsContextInternalServerError = new CosmosDiagnosticsContextCore();

            diagnosticsContextInternalServerError.AddDiagnosticsInternal(new PointOperationStatistics(
                                                                             Guid.NewGuid().ToString(),
                                                                             System.Net.HttpStatusCode.InternalServerError,
                                                                             subStatusCode: SubStatusCodes.Unknown,
                                                                             responseTimeUtc: DateTime.UtcNow,
                                                                             requestCharge: 10.2,
                                                                             errorMessage: "Error message",
                                                                             method: HttpMethod.Post,
                                                                             requestUri: new Uri("http://localhost.com"),
                                                                             requestSessionToken: null,
                                                                             responseSessionToken: null));

            diagnostics = new QueryPageDiagnostics(
                clientQueryCorrelationId: Guid.NewGuid(),
                partitionKeyRangeId: "0",
                queryMetricText: null,
                indexUtilizationText: null,
                diagnosticsContext: diagnosticsContextInternalServerError);
            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>(),
                                       cancellationTokenSource.Token)).Returns(
                Task.FromResult(QueryResponseCore.CreateFailure(
                                    statusCode: HttpStatusCode.InternalServerError,
                                    subStatusCodes: null,
                                    cosmosException: CosmosExceptionFactory.CreateInternalServerErrorException(
                                        "Error message"),
                                    requestCharge: 10.2,
                                    activityId: Guid.NewGuid().ToString())));

            await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);

            // First item should be a success
            {
                (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);

                Assert.IsTrue(movedToNextPage);
                Assert.IsNull(failureResponse);
                Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage());
                Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage());
                Assert.IsTrue(itemProducerTree.HasMoreResults);
            }

            // Second item should be a success
            {
                (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);

                Assert.IsTrue(movedToNextPage);
                Assert.IsNull(failureResponse);
                Assert.IsTrue(itemProducerTree.TryMoveNextDocumentWithinPage());
                Assert.IsFalse(itemProducerTree.TryMoveNextDocumentWithinPage());
                Assert.IsTrue(itemProducerTree.HasMoreResults);
            }

            // Third item should be a failure
            {
                (bool movedToNextPage, QueryResponseCore? failureResponse) = await itemProducerTree.TryMoveNextPageAsync(cancellationTokenSource.Token);

                Assert.IsFalse(movedToNextPage);
                Assert.IsNotNull(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>(),
                                       cancellationTokenSource.Token)).
            Throws(new Exception("Previous buffer failed. Operation should return original failure and not try again"));

            await itemProducerTree.BufferMoreDocumentsAsync(cancellationTokenSource.Token);

            Assert.IsFalse(itemProducerTree.HasMoreResults);
        }