From dafa6b811e27962c5feaa77c1807f817f4556e84 Mon Sep 17 00:00:00 2001 From: Sourabh Jain Date: Sun, 1 Sep 2024 07:13:16 +0530 Subject: [PATCH] fix tests --- .../src/Resource/ClientContextCore.cs | 109 +++++++++--------- .../OpenTelemetryRecorderFactory.cs | 4 +- .../OpenTelemetry/OpenTelemetryResponse.cs | 1 - .../Batch/BatchAsyncBatcherTests.cs | 4 +- .../Batch/BatchAsyncContainerExecutorTests.cs | 4 +- .../Batch/BatchAsyncStreamerTests.cs | 4 +- .../ChangeFeedEstimatorIteratorTests.cs | 8 +- ...dPartitionKeyResultSetIteratorCoreTests.cs | 16 +-- .../CosmosItemUnitTests.cs | 12 +- .../FeedRange/ChangeFeedIteratorCoreTests.cs | 6 +- 10 files changed, 87 insertions(+), 81 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs b/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs index 0c3b90e19f..45df0388dc 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs @@ -501,65 +501,71 @@ private async Task RunWithDiagnosticsHelperAsync( RequestOptions requestOptions, ResourceType? resourceType = null) { - if (resourceType is not null && this.IsBulkOperationSupported(resourceType.Value, operationType)) + using (OpenTelemetryCoreRecorder recorder = + OpenTelemetryRecorderFactory.CreateRecorder( + getOperationName: () => + { + if (resourceType is not null && this.IsBulkOperationSupported(resourceType.Value, operationType)) + { + operationName = OpenTelemetryConstants.Operations.ExecuteBulkPrefix + operationName; + } + + return operationName; + }, + containerName: containerName, + databaseName: databaseName, + operationType: operationType, + requestOptions: requestOptions, + trace: trace, + clientContext: this.isDisposed ? null : this)) { - operationName = OpenTelemetryConstants.Operations.ExecuteBulkPrefix + operationName; - } - - using (OpenTelemetryCoreRecorder recorder = - OpenTelemetryRecorderFactory.CreateRecorder( - operationName: operationName, - containerName: containerName, - databaseName: databaseName, - operationType: operationType, - requestOptions: requestOptions, - trace: trace, - clientContext: this.isDisposed ? null : this)) - using (new ActivityScope(Guid.NewGuid())) - { - try + using (new ActivityScope(Guid.NewGuid())) { - TResult result = await task(trace).ConfigureAwait(false); - if (openTelemetry != null && recorder.IsEnabled) + try { - // Record request response information - OpenTelemetryAttributes response = openTelemetry(result); - recorder.Record(response); + TResult result = await task(trace).ConfigureAwait(false); + if (openTelemetry != null && recorder.IsEnabled) + { + // Record request response information + OpenTelemetryAttributes response = openTelemetry(result); + recorder.Record(response); + } + + return result; } + catch (OperationCanceledException oe) when (!(oe is CosmosOperationCanceledException)) + { + CosmosOperationCanceledException operationCancelledException = new CosmosOperationCanceledException(oe, trace); + recorder.MarkFailed(operationCancelledException); - return result; - } - catch (OperationCanceledException oe) when (!(oe is CosmosOperationCanceledException)) - { - CosmosOperationCanceledException operationCancelledException = new CosmosOperationCanceledException(oe, trace); - recorder.MarkFailed(operationCancelledException); - - throw operationCancelledException; - } - catch (ObjectDisposedException objectDisposed) when (!(objectDisposed is CosmosObjectDisposedException)) - { - CosmosObjectDisposedException objectDisposedException = new CosmosObjectDisposedException( - objectDisposed, - this.client, - trace); - recorder.MarkFailed(objectDisposedException); + throw operationCancelledException; + } + catch (ObjectDisposedException objectDisposed) when (!(objectDisposed is CosmosObjectDisposedException)) + { + CosmosObjectDisposedException objectDisposedException = new CosmosObjectDisposedException( + objectDisposed, + this.client, + trace); + recorder.MarkFailed(objectDisposedException); - throw objectDisposedException; - } - catch (NullReferenceException nullRefException) when (!(nullRefException is CosmosNullReferenceException)) - { - CosmosNullReferenceException nullException = new CosmosNullReferenceException( - nullRefException, - trace); - recorder.MarkFailed(nullException); + throw objectDisposedException; + } + catch (NullReferenceException nullRefException) when (!(nullRefException is CosmosNullReferenceException)) + { + CosmosNullReferenceException nullException = new CosmosNullReferenceException( + nullRefException, + trace); + recorder.MarkFailed(nullException); - throw nullException; - } - catch (Exception ex) - { - recorder.MarkFailed(ex); + throw nullException; + } + catch (Exception ex) + { + recorder.MarkFailed(ex); + + throw; + } - throw; } } } @@ -585,7 +591,6 @@ private async Task ProcessResourceOperationAsBulkStreamAsync( resourceStream: streamPayload, requestOptions: batchItemRequestOptions, cosmosClientContext: this); - TransactionalBatchOperationResult batchOperationResult = await cosmosContainerCore.BatchExecutor.AddAsync( itemBatchOperation, trace, diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs index 1b16f48307..368594114a 100644 --- a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs @@ -26,7 +26,7 @@ internal static class OpenTelemetryRecorderFactory isStable: false), isThreadSafe: true); - public static OpenTelemetryCoreRecorder CreateRecorder(string operationName, + public static OpenTelemetryCoreRecorder CreateRecorder(Func getOperationName, string containerName, string databaseName, Documents.OperationType operationType, @@ -37,6 +37,8 @@ public static OpenTelemetryCoreRecorder CreateRecorder(string operationName, OpenTelemetryCoreRecorder openTelemetryRecorder = default; if (clientContext is { ClientOptions.CosmosClientTelemetryOptions.DisableDistributedTracing: false }) { + string operationName = getOperationName(); + // If there is no source then it will return default otherwise a valid diagnostic scope DiagnosticScope scope = LazyScopeFactory.Value.CreateScope(name: $"{OpenTelemetryAttributeKeys.OperationPrefix}.{operationName}", kind: clientContext.ClientOptions.ConnectionMode == ConnectionMode.Gateway ? ActivityKind.Internal : ActivityKind.Client); diff --git a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs index 5166940883..5738d58c5c 100644 --- a/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs +++ b/Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs @@ -8,7 +8,6 @@ namespace Microsoft.Azure.Cosmos using System.IO; using System.Net; using Microsoft.Azure.Cosmos.Core.Trace; - using Microsoft.Azure.Documents; using Telemetry; internal sealed class OpenTelemetryResponse : OpenTelemetryAttributes diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs index 5861121dd4..592590c83f 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs @@ -763,10 +763,10 @@ private static CosmosClientContext MockClientContext() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>, Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType,requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); return mockContext.Object; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs index 30c48043d5..506198a744 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs @@ -351,10 +351,10 @@ private Mock MockClientContext() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>, Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( (operationName,containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); mockContext.Setup(x => x.Client).Returns(MockCosmosUtil.CreateMockCosmosClient()); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs index 90db4644ef..641567d0a7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs @@ -207,10 +207,10 @@ private CosmosClientContext GetMockClientContext() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>, Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); return mockContext.Object; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedEstimatorIteratorTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedEstimatorIteratorTests.cs index 663017f214..41e14ea85b 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedEstimatorIteratorTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedEstimatorIteratorTests.cs @@ -392,10 +392,10 @@ static FeedIteratorInternal feedCreator(DocumentServiceLease lease, string conti It.IsAny(), It.IsAny>>>(), It.IsAny, OpenTelemetryAttributes>>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>>, Func, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>( + .Returns>>, Func, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) @@ -504,10 +504,10 @@ private static ContainerInternal GetMockedContainer() It.IsAny(), It.IsAny>>>(), It.IsAny, OpenTelemetryAttributes>>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>>, Func, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>( + .Returns>>, Func, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCoreTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCoreTests.cs index f5fdbf2831..8e90431099 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCoreTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCoreTests.cs @@ -49,10 +49,10 @@ public async Task EtagPassesContinuation() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>,Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>,Func, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) @@ -130,10 +130,10 @@ public async Task NextReadHasUpdatedContinuation() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>, Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) @@ -201,10 +201,10 @@ public async Task ShouldSetFeedRangePartitionKeyRange() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>, Func, ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) @@ -284,10 +284,10 @@ public async Task ShouldUseFeedRangeEpk() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.Is(tc => tc == TraceComponent.ChangeFeed), It.IsAny())) - .Returns>, Func, Documents.ResourceType, TraceComponent, TraceLevel>( + .Returns>, Func, Documents.ResourceType?, TraceComponent, TraceLevel>( (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => { using (ITrace trace = Trace.GetRootTrace(operationName, comp, level)) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs index f2ec22382d..1801104a7d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs @@ -717,11 +717,11 @@ public async Task TestMultipleNestedPartitionKeyValueFromStreamAsync() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>, Func, TraceComponent, TraceLevel>( - (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, comp, level) => func(NoOpTrace.Singleton)); + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( + (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); mockContext.Setup(x => x.OperationHelperAsync>( It.IsAny(), @@ -731,11 +731,11 @@ public async Task TestMultipleNestedPartitionKeyValueFromStreamAsync() It.IsAny(), It.IsAny>>>(), It.IsAny, OpenTelemetryAttributes>>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>>, Func, OpenTelemetryAttributes>, TraceComponent, TraceLevel>( - (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, comp, level) => func(NoOpTrace.Singleton)); + .Returns>>, Func, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>( + (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); mockContext.Setup(x => x.ProcessResourceOperationStreamAsync( It.IsAny(), diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ChangeFeedIteratorCoreTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ChangeFeedIteratorCoreTests.cs index 3d02743af9..4159da4080 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ChangeFeedIteratorCoreTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/FeedRange/ChangeFeedIteratorCoreTests.cs @@ -483,11 +483,11 @@ private CosmosClientContext MockClientContext() It.IsAny(), It.IsAny>>(), It.IsAny>(), - It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny())) - .Returns>, Func, TraceComponent, TraceLevel>( - (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, comp, level) => func(NoOpTrace.Singleton)); + .Returns>, Func, ResourceType?, TraceComponent, TraceLevel>( + (operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton)); return mockContext.Object; }