CosmosOperationCanceledExceptionTests.cs

//------------------------------------------------------------
// Copyright (c) Microsoft Corporation.  All rights reserved.
//------------------------------------------------------------

namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
{
    using System;
    using System.Threading;
    using System.Threading.Tasks;
    using Microsoft.Azure.Cosmos.Diagnostics;
    using Microsoft.Azure.Cosmos.Tracing;
    using Microsoft.VisualStudio.TestTools.UnitTesting;

    [TestClass]
    public class CosmosOperationCanceledExceptionTests : BaseCosmosClientHelper
    {
        private ContainerInternal Container = null;

        [TestInitialize]
        public async Task TestInitialize()
        {
            await base.TestInit();
            string PartitionKey = "/pk";
            ContainerResponse response = await this.database.CreateContainerAsync(
                new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: PartitionKey),
                cancellationToken: this.cancellationToken);
            Assert.IsNotNull(response);
            Assert.IsNotNull(response.Container);
            Assert.IsNotNull(response.Resource);
            this.Container = (ContainerInternal)response;
        }

        [TestCleanup]
        public async Task Cleanup()
        {
            await base.TestCleanup();
        }

        [TestMethod]
        public async Task CheckCancellationTokenGatewayTestAsync()
        {
            using (CosmosClient gatewayClient = TestCommon.CreateCosmosClient(
                builder => builder.WithConnectionModeGateway()))
            {
                Container gatewayContainer = gatewayClient.GetContainer(this.database.Id, this.Container.Id);
                CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
                cancellationTokenSource.Cancel();
                await this.CheckCancellationTokenTestAsync(gatewayContainer, cancellationTokenSource.Token);
            }
        }

        [TestMethod]
        public async Task CheckCancellationWithTransportIntercepterTestAsync()
        {
            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
            Container withCancellationToken = TransportClientHelper.GetContainerWithIntercepter(
                 this.database.Id,
                 this.Container.Id,
                 (uri, resourceOperation, documentServiceRequest) =>
                 {
                     if (documentServiceRequest.ResourceType == Documents.ResourceType.Document)
                     {
                         cancellationTokenSource.Cancel();
                     }
                 },
                 useGatewayMode: false,
                 (uri, resourceOperation, documentServiceRequest) 
                    => TransportClientHelper.ReturnThrottledStoreResponseOnItemOperation(uri, resourceOperation, documentServiceRequest, Guid.NewGuid(), string.Empty));

            await this.CheckCancellationTokenTestAsync(withCancellationToken, cancellationTokenSource.Token);
        }

        [TestMethod]
        public async Task CheckCancellationTokenDirectTestAsync()
        {
            CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
            cancellationTokenSource.Cancel();
            await this.CheckCancellationTokenTestAsync(this.Container, cancellationTokenSource.Token);
        }

        [TestMethod]
        public void CheckToJsonStringConversion()
        {
            // This test is reproducing an issue that an internal customer faced after enabling
            // cross region hedging 
            // In their scenario they occassionally hit e2e timeouts as CosmosOperationCancelledException
            // The JSON transformation of the CosmosOperationCancelledException in this case could run into
            // an issue (InvalidOperationException due to Lazy<T> factory hitting a cyclic dependency
            // when neither the diagnostics nor the exception had materialized the json string yet
            // because the CosmosOperationCancelledException ias added as a trace datum to the ITrace of
            // the CosmosTraceDiagnostics instance.
            // The test below reproduced the issue - and is kept here to validate the fix and
            // as a regression test
            using ITrace outerTrace = Trace.GetRootTrace("cyclicDependencyReproOuter");
            using ITrace innerTrace = outerTrace.StartChild("cyclicDependencyReproInner");

            OperationCanceledException innerTimeout = new OperationCanceledException();
            CosmosOperationCanceledException innerCosmosTimeoutException = new CosmosOperationCanceledException(
                innerTimeout,
                innerTrace);
            CosmosOperationCanceledException outerCosmosTimeoutException = new CosmosOperationCanceledException(
                innerCosmosTimeoutException,
                outerTrace);

            Console.WriteLine(outerCosmosTimeoutException.ToString());
        }

        private async Task CheckCancellationTokenTestAsync(
            Container container,
            CancellationToken cancellationToken)
        {
            string diagnosticsMessageStartString = "CosmosDiagnostics: {\"Summary\":{";
            ToDoActivity toDoActivity = ToDoActivity.CreateRandomToDoActivity();

            try
            {
                await container.CreateItemAsync<ToDoActivity>(
                    toDoActivity,
                    new Cosmos.PartitionKey(toDoActivity.pk),
                    cancellationToken: cancellationToken);

                Assert.Fail("Should have thrown");
            }
            catch (CosmosOperationCanceledException ce)
            {
                Assert.IsNotNull(ce);
                string message = ce.Message;
                string diagnostics = ce.Diagnostics.ToString();
                string toString = ce.ToString();
                Assert.IsTrue(toString.Contains(diagnosticsMessageStartString), $"Exception ToString() : {toString} does not contain diagnostics: {diagnostics}");
                Assert.IsTrue(message.Contains(diagnosticsMessageStartString), $"Message : {message} does not contain diagnostics: {diagnostics}");
                string messageWithoutDiagnostics = message[..message.IndexOf(Environment.NewLine)].Trim();
                Assert.IsTrue(toString.Contains(messageWithoutDiagnostics), $"Exception ToString() : {toString} does not contain message: {messageWithoutDiagnostics}");
            }

            try
            {
                FeedIterator feedIterator = container.GetItemQueryStreamIterator(
                    "select * from T");

                await feedIterator.ReadNextAsync(cancellationToken: cancellationToken);

                Assert.Fail("Should have thrown");
            }
            catch (CosmosOperationCanceledException ce)
            {
                string message = ce.Message;
                string diagnostics = ce.Diagnostics.ToString();
                string toString = ce.ToString();
                Assert.IsTrue(toString.Contains(diagnosticsMessageStartString), $"Exception ToString() : {toString} does not contain diagnostics: {diagnostics}");
                Assert.IsTrue(message.Contains(diagnosticsMessageStartString), $"Message : {message} does not contain diagnostics: {diagnostics}");
                string messageWithoutDiagnostics = message[..message.IndexOf(Environment.NewLine)].Trim();
                Assert.IsTrue(toString.Contains(messageWithoutDiagnostics), $"Exception ToString() : {toString} does not contain message: {messageWithoutDiagnostics}");
            }
        }
    }
}