Skip to content

Commit 654a6cd

Browse files
alexsohn1126alexsohn1126
committed
cleanup tests and other stuff
1 parent e612c00 commit 654a6cd

File tree

5 files changed

+129
-65
lines changed

5 files changed

+129
-65
lines changed

‎src/Sentry.Extensions.AI/SentryAISpanEnricher.cs‎

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ public static void EnrichWithStreamingResponses(ISpan span, List<ChatResponseUpd
121121
}
122122
}
123123

124-
if (aiOptions?.IncludeAIResponseContent == true)
124+
if (aiOptions.IncludeAIResponseContent)
125125
{
126126
span.SetData(SentryAIConstants.SpanAttributes.ResponseText, finalText.ToString());
127127
}
@@ -141,19 +141,39 @@ private static void PopulateToolCallsInfo(IList<AIContent> contents, ISpan span)
141141
}
142142
}
143143

144-
private static string FormatAvailableTools(IList<AITool> tools) =>
145-
FormatAsJson(tools, tool => new
144+
private static string FormatAvailableTools(IList<AITool> tools)
145+
{
146+
try
146147
{
147-
name = tool.Name,
148-
description = tool.Description
149-
});
148+
var str = FormatAsJson(tools, tool => new
149+
{
150+
name = tool.Name,
151+
description = tool.Description
152+
});
153+
return str;
154+
}
155+
catch
156+
{
157+
return "";
158+
}
159+
}
150160

151-
private static string FormatRequestMessage(ChatMessage[] messages) =>
152-
FormatAsJson(messages, message => new
161+
private static string FormatRequestMessage(ChatMessage[] messages)
162+
{
163+
try
153164
{
154-
role = message.Role,
155-
content = message.Text
156-
});
165+
var str = FormatAsJson(messages, message => new
166+
{
167+
role = message.Role,
168+
content = message.Text
169+
});
170+
return str;
171+
}
172+
catch
173+
{
174+
return "";
175+
}
176+
}
157177

158178
private static string FormatFunctionCallContent(FunctionCallContent[] content) =>
159179
FormatAsJson(content, c => new

‎src/Sentry.Extensions.AI/SentryChatClient.cs‎

Lines changed: 11 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,18 @@ internal sealed class SentryChatClient : DelegatingChatClient
1212
public SentryChatClient(IChatClient client, Action<SentryAIOptions>? configure = null) : base(client)
1313
{
1414
_sentryAIOptions = new SentryAIOptions();
15-
configure?.Invoke(_sentryAIOptions);
15+
try
16+
{
17+
configure?.Invoke(_sentryAIOptions);
1618

17-
if (_sentryAIOptions.InitializeSdk && !SentrySdk.IsEnabled)
19+
if (_sentryAIOptions.InitializeSdk && !SentrySdk.IsEnabled)
20+
{
21+
SentrySdk.Init(_sentryAIOptions);
22+
}
23+
}
24+
catch (Exception ex)
1825
{
19-
SentrySdk.Init(_sentryAIOptions);
26+
SentrySdk.CaptureException(ex);
2027
}
2128
}
2229

@@ -47,15 +54,6 @@ public override async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessag
4754
_hub.CaptureException(ex);
4855
throw;
4956
}
50-
finally
51-
{
52-
// if options was null, we need to finish root span immediately because no tool calls will be made
53-
// therefore no consequent GetResponseAsync calls
54-
if (options == null)
55-
{
56-
outerSpan?.Finish();
57-
}
58-
}
5957
}
6058

6159
/// <inheritdoc cref="IChatClient"/>
@@ -69,7 +67,6 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
6967
var outerSpan = TryGetRootSpan(options);
7068
var innerSpan = CreateChatSpan(outerSpan, options);
7169

72-
var hasNext = false;
7370
var responses = new List<ChatResponseUpdate>();
7471
var enumerator = base
7572
.GetStreamingResponseAsync(chatMessages, options, cancellationToken)
@@ -81,7 +78,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
8178
ChatResponseUpdate? current;
8279
try
8380
{
84-
hasNext = await enumerator.MoveNextAsync().ConfigureAwait(false);
81+
var hasNext = await enumerator.MoveNextAsync().ConfigureAwait(false);
8582

8683
if (!hasNext)
8784
{
@@ -100,14 +97,6 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
10097
_hub.CaptureException(ex);
10198
throw;
10299
}
103-
finally
104-
{
105-
// if options was null, and we don't have next text, we need to finish root span immediately
106-
if (options == null && !hasNext)
107-
{
108-
outerSpan?.Finish(SpanStatus.Ok);
109-
}
110-
}
111100

112101
yield return current;
113102
}

‎test/Sentry.Extensions.AI.Tests/SentryAIActivityListenerTests.cs‎

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -5,28 +5,6 @@ namespace Sentry.Extensions.AI.Tests;
55

66
public class SentryAIActivityListenerTests
77
{
8-
private class Fixture
9-
{
10-
private SentryOptions Options { get; }
11-
public ISentryClient Client { get; }
12-
public IHub Hub { get; set; }
13-
14-
public Fixture()
15-
{
16-
Options = new SentryOptions
17-
{
18-
Dsn = ValidDsn,
19-
TracesSampleRate = 1.0,
20-
};
21-
22-
Hub = Substitute.For<IHub>();
23-
Client = Substitute.For<ISentryClient>();
24-
SentrySdk.Init(Options);
25-
}
26-
}
27-
28-
private readonly Fixture _fixture = new();
29-
308
[Fact]
319
public void Init_AddsActivityListenerToActivitySource()
3210
{

‎test/Sentry.Extensions.AI.Tests/SentryChatClientTests.cs‎

Lines changed: 86 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,53 +5,131 @@ namespace Sentry.Extensions.AI.Tests;
55

66
public class SentryChatClientTests
77
{
8+
private class Fixture
9+
{
10+
private SentryOptions Options { get; }
11+
public ISentryClient Client { get; }
12+
public IHub Hub { get; set; }
13+
14+
public Fixture()
15+
{
16+
Options = new SentryOptions
17+
{
18+
Dsn = ValidDsn,
19+
TracesSampleRate = 1.0,
20+
};
21+
22+
SentrySdk.Init(Options);
23+
Hub = SentrySdk.CurrentHub;
24+
Client = Substitute.For<ISentryClient>();
25+
}
26+
}
27+
28+
private readonly Fixture _fixture = new();
29+
830
[Fact]
931
public async Task CompleteAsync_CallsInnerClient()
1032
{
33+
// Arrange
1134
var inner = Substitute.For<IChatClient>();
35+
var sentryChatClient = new SentryChatClient(inner);
1236
var message = new ChatMessage(ChatRole.Assistant, "ok");
1337
var chatResponse = new ChatResponse(message);
1438
inner.GetResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(), Arg.Any<CancellationToken>())
1539
.Returns(Task.FromResult(chatResponse));
1640

17-
var sentryChatClient = new SentryChatClient(inner);
18-
41+
// Act
1942
var res = await sentryChatClient.GetResponseAsync([new ChatMessage(ChatRole.User, "hi")]);
2043

44+
// Assert
2145
Assert.Equal([message], res.Messages);
2246
await inner.Received(1).GetResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(),
2347
Arg.Any<CancellationToken>());
2448
}
2549

2650
[Fact]
27-
public async Task CompleteStreamingAsync_CallsInnerClient()
51+
public async Task CompleteStreamingAsync_CallsInnerClient_AndSetsSpanData()
2852
{
29-
var inner = Substitute.For<IChatClient>();
53+
// Arrange - Use Fixture Hub to start transaction
54+
var transaction = _fixture.Hub.StartTransaction("test-streaming", "test");
3055

56+
_fixture.Hub.ConfigureScope(scope => scope.Transaction = transaction);
57+
SentrySdk.ConfigureScope(scope => scope.Transaction = transaction);
58+
59+
var inner = Substitute.For<IChatClient>();
3160
inner.GetStreamingResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(),
3261
Arg.Any<CancellationToken>())
3362
.Returns(CreateTestStreamingUpdatesAsync());
34-
3563
var client = new SentryChatClient(inner);
36-
3764
var results = new List<ChatResponseUpdate>();
65+
66+
// Act
3867
await foreach (var update in client.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "hi")]))
3968
{
4069
results.Add(update);
4170
}
4271

72+
// Assert
4373
Assert.Equal(2, results.Count);
4474
Assert.Equal("Hello", results[0].Text);
4575
Assert.Equal(" World!", results[1].Text);
46-
4776
inner.Received(1).GetStreamingResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(),
4877
Arg.Any<CancellationToken>());
78+
var spans = transaction.Spans;
79+
var chatSpan = spans.FirstOrDefault(s => s.Operation == SentryAIConstants.SpanAttributes.ChatOperation);
80+
Assert.NotNull(chatSpan);
81+
Assert.Equal(SpanStatus.Ok, chatSpan.Status);
82+
Assert.True(chatSpan.IsFinished);
83+
Assert.Equal("chat", chatSpan.Data[SentryAIConstants.SpanAttributes.OperationName]);
84+
Assert.Equal("Hello World!", chatSpan.Data[SentryAIConstants.SpanAttributes.ResponseText]);
85+
}
86+
87+
[Fact]
88+
public async Task CompleteStreamingAsync_HandlesErrors_AndFinishesSpanWithException()
89+
{
90+
// Arrange
91+
var transaction = _fixture.Hub.StartTransaction("test-streaming-error", "test");
92+
_fixture.Hub.ConfigureScope(scope => scope.Transaction = transaction);
93+
SentrySdk.ConfigureScope(scope => scope.Transaction = transaction);
94+
95+
var inner = Substitute.For<IChatClient>();
96+
var expectedException = new InvalidOperationException("Streaming failed");
97+
inner.GetStreamingResponseAsync(Arg.Any<IList<ChatMessage>>(), Arg.Any<ChatOptions>(),
98+
Arg.Any<CancellationToken>())
99+
.Returns(CreateFailingStreamingUpdatesAsync(expectedException));
100+
var client = new SentryChatClient(inner);
101+
102+
// Act
103+
var actualException = await Assert.ThrowsAsync<InvalidOperationException>(async () =>
104+
{
105+
await foreach (var update in client.GetStreamingResponseAsync([new ChatMessage(ChatRole.User, "hi")]))
106+
{
107+
// Should not reach here due to exception
108+
}
109+
});
110+
111+
Assert.Equal(expectedException.Message, actualException.Message);
112+
113+
// Assert
114+
var spans = transaction.Spans;
115+
var chatSpan = spans.FirstOrDefault(s => s.Operation == SentryAIConstants.SpanAttributes.ChatOperation);
116+
Assert.NotNull(chatSpan);
117+
Assert.Equal(SpanStatus.InternalError, chatSpan.Status);
118+
Assert.True(chatSpan.IsFinished);
119+
Assert.Equal("chat", chatSpan.Data[SentryAIConstants.SpanAttributes.OperationName]);
120+
}
121+
122+
private static async IAsyncEnumerable<ChatResponseUpdate> CreateFailingStreamingUpdatesAsync(Exception exception)
123+
{
124+
yield return new ChatResponseUpdate(ChatRole.System, "Hello");
125+
await Task.Yield();
126+
throw exception;
49127
}
50128

51129
private static async IAsyncEnumerable<ChatResponseUpdate> CreateTestStreamingUpdatesAsync()
52130
{
53131
yield return new ChatResponseUpdate(ChatRole.System, "Hello");
54-
await Task.Yield(); // Make it async
132+
await Task.Yield();
55133
yield return new ChatResponseUpdate(ChatRole.System, " World!");
56134
}
57135
}

‎test/Sentry.Extensions.AI.Tests/SentryInstrumentedFunctionTests.cs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,6 @@ public async Task InvokeCoreAsync_WithNullResult_ReturnsNull()
4747
var result = await sentryFunction.InvokeAsync(arguments);
4848

4949
// Assert
50-
// AIFunctionFactory may return JsonElement with ValueKind.Null instead of actual null
5150
if (result is JsonElement jsonElement)
5251
{
5352
Assert.Equal(JsonValueKind.Null, jsonElement.ValueKind);
@@ -219,7 +218,7 @@ public static IDisposable InitializeSdk()
219218
{
220219
return SentrySdk.Init(options =>
221220
{
222-
options.Dsn = "https://[email protected]:65535/2147483647";
221+
options.Dsn = ValidDsn;
223222
options.TracesSampleRate = 1.0;
224223
options.Debug = false;
225224
});

0 commit comments

Comments
 (0)