Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@sourabh1007
sourabh1007 committed Sep 1, 2024
1 parent 717ca49 commit dafa6b8
Show file tree
Hide file tree
Showing 10 changed files with 87 additions and 81 deletions.
109 changes: 57 additions & 52 deletions Microsoft.Azure.Cosmos/src/Resource/ClientContextCore.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -501,65 +501,71 @@ private async Task<TResult> RunWithDiagnosticsHelperAsync<TResult>(
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;
}
}
}
Expand All @@ -585,7 +591,6 @@ private async Task<ResponseMessage> ProcessResourceOperationAsBulkStreamAsync(
resourceStream: streamPayload,
requestOptions: batchItemRequestOptions,
cosmosClientContext: this);

TransactionalBatchOperationResult batchOperationResult = await cosmosContainerCore.BatchExecutor.AddAsync(
itemBatchOperation,
trace,
Expand Down
4 changes: 3 additions & 1 deletion Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ internal static class OpenTelemetryRecorderFactory
isStable: false),
isThreadSafe: true);

public static OpenTelemetryCoreRecorder CreateRecorder(string operationName,
public static OpenTelemetryCoreRecorder CreateRecorder(Func<string> getOperationName,
string containerName,
string databaseName,
Documents.OperationType operationType,
Expand All @@ -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);
Expand Down
1 change: 0 additions & 1 deletion Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncBatcherTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -763,10 +763,10 @@ private static CosmosClientContext MockClientContext()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<object>>>(),
It.IsAny<Func<object, OpenTelemetryAttributes>>(),
It.IsAny<ResourceType>(),
It.IsAny<ResourceType?>(),
It.IsAny<TraceComponent>(),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType,requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton));

return mockContext.Object;
Expand Down
4 changes: 2 additions & 2 deletions ...Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncContainerExecutorTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -351,10 +351,10 @@ private Mock<CosmosClientContext> MockClientContext()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<object>>>(),
It.IsAny<Func<object, OpenTelemetryAttributes>>(),
It.IsAny<ResourceType>(),
It.IsAny<ResourceType?>(),
It.IsAny<TraceComponent>(),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName,containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton));

mockContext.Setup(x => x.Client).Returns(MockCosmosUtil.CreateMockCosmosClient());
Expand Down
4 changes: 2 additions & 2 deletions Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,10 +207,10 @@ private CosmosClientContext GetMockClientContext()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<object>>>(),
It.IsAny<Func<object, OpenTelemetryAttributes>>(),
It.IsAny<ResourceType>(),
It.IsAny<ResourceType?>(),
It.IsAny<TraceComponent>(),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, OperationType, RequestOptions, Func<ITrace, Task<object>>, Func<object, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) => func(NoOpTrace.Singleton));

return mockContext.Object;
Expand Down
8 changes: 4 additions & 4 deletions ....Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedEstimatorIteratorTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -392,10 +392,10 @@ static FeedIteratorInternal feedCreator(DocumentServiceLease lease, string conti
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>>(),
It.IsAny<Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>>(),
It.IsAny<ResourceType>(),
It.IsAny<ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>, Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>, Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down Expand Up @@ -504,10 +504,10 @@ private static ContainerInternal GetMockedContainer()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>>(),
It.IsAny<Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>>(),
It.IsAny<ResourceType>(),
It.IsAny<ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>, Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<FeedResponse<ChangeFeedProcessorState>>>, Func<FeedResponse<ChangeFeedProcessorState>, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down
16 changes: 8 additions & 8 deletions ...crosoft.Azure.Cosmos.Tests/ChangeFeed/ChangeFeedPartitionKeyResultSetIteratorCoreTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,10 @@ public async Task EtagPassesContinuation()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<ResponseMessage>>>(),
It.IsAny<Func<ResponseMessage, OpenTelemetryAttributes>>(),
It.IsAny<Documents.ResourceType>(),
It.IsAny<Documents.ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>,Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>,Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down Expand Up @@ -130,10 +130,10 @@ public async Task NextReadHasUpdatedContinuation()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<ResponseMessage>>>(),
It.IsAny<Func<ResponseMessage, OpenTelemetryAttributes>>(),
It.IsAny<Documents.ResourceType>(),
It.IsAny<Documents.ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down Expand Up @@ -201,10 +201,10 @@ public async Task ShouldSetFeedRangePartitionKeyRange()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<ResponseMessage>>>(),
It.IsAny<Func<ResponseMessage, OpenTelemetryAttributes>>(),
It.IsAny<Documents.ResourceType>(),
It.IsAny<Documents.ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down Expand Up @@ -284,10 +284,10 @@ public async Task ShouldUseFeedRangeEpk()
It.IsAny<RequestOptions>(),
It.IsAny<Func<ITrace, Task<ResponseMessage>>>(),
It.IsAny<Func<ResponseMessage, OpenTelemetryAttributes>>(),
It.IsAny<Documents.ResourceType>(),
It.IsAny<Documents.ResourceType?>(),
It.Is<TraceComponent>(tc => tc == TraceComponent.ChangeFeed),
It.IsAny<TraceLevel>()))
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, Documents.ResourceType, TraceComponent, TraceLevel>(
.Returns<string, string, string, Documents.OperationType, RequestOptions, Func<ITrace, Task<ResponseMessage>>, Func<ResponseMessage, OpenTelemetryAttributes>, Documents.ResourceType?, TraceComponent, TraceLevel>(
(operationName, containerName, databaseName, operationType, requestOptions, func, oTelFunc, resourceType, comp, level) =>
{
using (ITrace trace = Trace.GetRootTrace(operationName, comp, level))
Expand Down
Loading

0 comments on commit dafa6b8

Please sign in to comment.