Skip to content

Commit 05fbd75

Browse files
reakaleekreakaleek
committed
Refactor and dispose disposable
1 parent 006b096 commit 05fbd75

File tree

6 files changed

+117
-192
lines changed

6 files changed

+117
-192
lines changed

‎src/api/Elastic.Documentation.Api.Infrastructure/Adapters/AskAi/AgentBuilderAskAiGateway.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,15 +37,15 @@ public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx = default)
3737
var kibanaUrl = await parameterProvider.GetParam("docs-kibana-url", false, ctx);
3838
var kibanaApiKey = await parameterProvider.GetParam("docs-kibana-apikey", true, ctx);
3939

40-
var request = new HttpRequestMessage(HttpMethod.Post,
40+
using var request = new HttpRequestMessage(HttpMethod.Post,
4141
$"{kibanaUrl}/api/agent_builder/converse/async")
4242
{
4343
Content = new StringContent(requestBody, Encoding.UTF8, "application/json")
4444
};
4545
request.Headers.Add("kbn-xsrf", "true");
4646
request.Headers.Authorization = new AuthenticationHeaderValue("ApiKey", kibanaApiKey);
4747

48-
var response = await httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, ctx);
48+
using var response = await httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, ctx);
4949

5050
// Ensure the response is successful before streaming
5151
if (!response.IsSuccessStatusCode)

‎src/api/Elastic.Documentation.Api.Infrastructure/Adapters/AskAi/LlmGatewayAskAiGateway.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx = default)
2525
{
2626
var llmGatewayRequest = LlmGatewayRequest.CreateFromRequest(askAiRequest);
2727
var requestBody = JsonSerializer.Serialize(llmGatewayRequest, LlmGatewayContext.Default.LlmGatewayRequest);
28-
var request = new HttpRequestMessage(HttpMethod.Post, options.FunctionUrl)
28+
using var request = new HttpRequestMessage(HttpMethod.Post, options.FunctionUrl)
2929
{
3030
Content = new StringContent(requestBody, Encoding.UTF8, "application/json")
3131
};
@@ -37,7 +37,7 @@ public async Task<Stream> AskAi(AskAiRequest askAiRequest, Cancel ctx = default)
3737

3838
// Use HttpCompletionOption.ResponseHeadersRead to get headers immediately
3939
// This allows us to start streaming as soon as headers are received
40-
var response = await httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, ctx);
40+
using var response = await httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, ctx);
4141

4242
// Ensure the response is successful before streaming
4343
if (!response.IsSuccessStatusCode)

‎tests-integration/Elastic.Documentation.Api.IntegrationTests/EuidEnrichmentIntegrationTests.cs‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,18 @@
77
using Elastic.Documentation.Api.Core;
88
using Elastic.Documentation.Api.Core.AskAi;
99
using Elastic.Documentation.Api.IntegrationTests.Fixtures;
10+
using FakeItEasy;
1011
using FluentAssertions;
12+
using Microsoft.Extensions.DependencyInjection;
1113

1214
namespace Elastic.Documentation.Api.IntegrationTests;
1315

1416
/// <summary>
1517
/// Integration tests for euid cookie enrichment in OpenTelemetry traces and logging.
16-
/// Uses WebApplicationFactory to test the real API configuration with mocked services.
18+
/// Uses WebApplicationFactory to test the real API configuration with mocked AskAi services.
1719
/// </summary>
18-
public class EuidEnrichmentIntegrationTests(ApiWebApplicationFactory factory) : IClassFixture<ApiWebApplicationFactory>
20+
public class EuidEnrichmentIntegrationTests
1921
{
20-
private readonly ApiWebApplicationFactory _factory = factory;
21-
2222
/// <summary>
2323
/// Test that verifies euid cookie is added to both HTTP span and custom AskAi span,
2424
/// and appears in log entries - using the real API configuration.
@@ -29,8 +29,39 @@ public async Task AskAiEndpointPropagatatesEuidToAllSpansAndLogs()
2929
// Arrange
3030
const string expectedEuid = "integration-test-euid-12345";
3131

32+
// Track streams created by mocks so we can dispose them after the test
33+
var mockStreams = new List<MemoryStream>();
34+
35+
// Create factory with mocked AskAi services
36+
using var factory = ApiWebApplicationFactory.WithMockedServices(services =>
37+
{
38+
// Mock IAskAiGateway to avoid external AI service calls
39+
var mockAskAiGateway = A.Fake<IAskAiGateway<Stream>>();
40+
A.CallTo(() => mockAskAiGateway.AskAi(A<AskAiRequest>._, A<Cancel>._))
41+
.ReturnsLazily(() =>
42+
{
43+
var stream = new MemoryStream(Encoding.UTF8.GetBytes("data: test\n\n"));
44+
mockStreams.Add(stream);
45+
return Task.FromResult<Stream>(stream);
46+
});
47+
services.AddSingleton(mockAskAiGateway);
48+
49+
// Mock IStreamTransformer
50+
var mockTransformer = A.Fake<IStreamTransformer>();
51+
A.CallTo(() => mockTransformer.AgentProvider).Returns("test-provider");
52+
A.CallTo(() => mockTransformer.AgentId).Returns("test-agent");
53+
A.CallTo(() => mockTransformer.TransformAsync(A<Stream>._, A<string?>._, A<Activity?>._, A<Cancel>._))
54+
.ReturnsLazily((Stream s, string? _, Activity? activity, Cancel _) =>
55+
{
56+
// Dispose the activity if provided (simulating what the real transformer does)
57+
activity?.Dispose();
58+
return Task.FromResult(s);
59+
});
60+
services.AddSingleton(mockTransformer);
61+
});
62+
3263
// Create client
33-
using var client = _factory.CreateClient();
64+
using var client = factory.CreateClient();
3465

3566
// Act - Make request to /ask-ai/stream with euid cookie
3667
using var request = new HttpRequestMessage(HttpMethod.Post, "/docs/_api/v1/ask-ai/stream");
@@ -48,7 +79,7 @@ public async Task AskAiEndpointPropagatatesEuidToAllSpansAndLogs()
4879
response.IsSuccessStatusCode.Should().BeTrue();
4980

5081
// Assert - Verify spans were captured
51-
var activities = _factory.ExportedActivities;
82+
var activities = factory.ExportedActivities;
5283
activities.Should().NotBeEmpty("OpenTelemetry should have captured activities");
5384

5485
// Verify HTTP span has euid
@@ -67,7 +98,7 @@ public async Task AskAiEndpointPropagatatesEuidToAllSpansAndLogs()
6798
askAiEuidTag.Value.Should().Be(expectedEuid, "AskAi span euid should match cookie value");
6899

69100
// Assert - Verify logs have euid in attributes
70-
var logRecords = _factory.ExportedLogRecords;
101+
var logRecords = factory.ExportedLogRecords;
71102
logRecords.Should().NotBeEmpty("Should have captured log records");
72103

73104
// Find a log entry from AskAiUsecase
@@ -80,5 +111,9 @@ public async Task AskAiEndpointPropagatatesEuidToAllSpansAndLogs()
80111
var euidAttribute = askAiLogRecord!.Attributes?.FirstOrDefault(a => a.Key == TelemetryConstants.UserEuidAttributeName) ?? default;
81112
euidAttribute.Should().NotBe(default(KeyValuePair<string, object?>), "Log record should include user.euid attribute");
82113
(euidAttribute.Value?.ToString() ?? string.Empty).Should().Be(expectedEuid, "Log record euid should match cookie value");
114+
115+
// Cleanup - dispose all mock streams
116+
foreach (var stream in mockStreams)
117+
stream.Dispose();
83118
}
84119
}

‎tests-integration/Elastic.Documentation.Api.IntegrationTests/Examples/ServiceMockingExampleTests.cs‎

Lines changed: 0 additions & 95 deletions
This file was deleted.

‎tests-integration/Elastic.Documentation.Api.IntegrationTests/Fixtures/ApiWebApplicationFactory.cs‎

Lines changed: 31 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@
33
// See the LICENSE file in the project root for more information
44

55
using System.Diagnostics;
6-
using System.Text;
7-
using Elastic.Documentation.Api.Core.AskAi;
8-
using Elastic.Documentation.Api.Core.Telemetry;
96
using Elastic.Documentation.Api.Infrastructure;
107
using Elastic.Documentation.Api.Infrastructure.Aws;
118
using Elastic.Documentation.Api.Infrastructure.OpenTelemetry;
@@ -14,7 +11,6 @@
1411
using Microsoft.AspNetCore.Mvc.Testing;
1512
using Microsoft.Extensions.DependencyInjection;
1613
using Microsoft.Extensions.DependencyInjection.Extensions;
17-
using Microsoft.Extensions.Logging;
1814
using OpenTelemetry;
1915
using OpenTelemetry.Logs;
2016
using OpenTelemetry.Trace;
@@ -24,12 +20,13 @@ namespace Elastic.Documentation.Api.IntegrationTests.Fixtures;
2420
/// <summary>
2521
/// Custom WebApplicationFactory for testing the API with mocked services.
2622
/// This fixture can be reused across multiple test classes.
23+
/// Only mocks services that ALL tests need (OpenTelemetry, AWS Parameters).
24+
/// Test-specific mocks should be configured using WithMockedServices.
2725
/// </summary>
2826
public class ApiWebApplicationFactory : WebApplicationFactory<Program>
2927
{
3028
public List<Activity> ExportedActivities { get; } = [];
3129
public List<LogRecord> ExportedLogRecords { get; } = [];
32-
private readonly List<MemoryStream> _mockMemoryStreams = [];
3330
private readonly Action<IServiceCollection>? _configureServices;
3431

3532
public ApiWebApplicationFactory() : this(null)
@@ -60,67 +57,31 @@ public static ApiWebApplicationFactory WithMockedServices(Action<IServiceCollect
6057
=> new(configureServices);
6158

6259
protected override void ConfigureWebHost(IWebHostBuilder builder) => builder.ConfigureServices(services =>
63-
{
64-
var otelBuilder = services.AddOpenTelemetry();
65-
_ = otelBuilder.WithTracing(tracing =>
66-
{
67-
_ = tracing
68-
.AddDocsApiTracing() // Reuses production configuration
69-
.AddInMemoryExporter(ExportedActivities);
70-
});
71-
_ = otelBuilder.WithLogging(logging =>
72-
{
73-
_ = logging
74-
.AddDocsApiLogging() // Reuses production configuration
75-
.AddInMemoryExporter(ExportedLogRecords);
76-
});
77-
78-
// Mock IParameterProvider to avoid AWS dependencies
79-
var mockParameterProvider = A.Fake<IParameterProvider>();
80-
A.CallTo(() => mockParameterProvider.GetParam(A<string>._, A<bool>._, A<Cancel>._))
81-
.Returns(Task.FromResult("mock-value"));
82-
_ = services.AddSingleton(mockParameterProvider);
83-
84-
// Mock IAskAiGateway to avoid external AI service calls
85-
var mockAskAiGateway = A.Fake<IAskAiGateway<Stream>>();
86-
A.CallTo(() => mockAskAiGateway.AskAi(A<AskAiRequest>._, A<Cancel>._))
87-
.ReturnsLazily(() =>
88-
{
89-
var stream = new MemoryStream(Encoding.UTF8.GetBytes("data: test\n\n"));
90-
_mockMemoryStreams.Add(stream);
91-
return Task.FromResult<Stream>(stream);
92-
});
93-
_ = services.AddSingleton(mockAskAiGateway);
94-
95-
// Mock IStreamTransformer
96-
var mockTransformer = A.Fake<IStreamTransformer>();
97-
A.CallTo(() => mockTransformer.AgentProvider).Returns("test-provider");
98-
A.CallTo(() => mockTransformer.AgentId).Returns("test-agent");
99-
A.CallTo(() => mockTransformer.TransformAsync(A<Stream>._, A<string?>._, A<Activity?>._, A<Cancel>._))
100-
.ReturnsLazily((Stream s, string? _, Activity? activity, Cancel _) =>
101-
{
102-
// Dispose the activity if provided (simulating what the real transformer does)
103-
activity?.Dispose();
104-
return Task.FromResult(s);
105-
});
106-
_ = services.AddSingleton(mockTransformer);
107-
108-
// Allow tests to override services - RemoveAll + Add to properly replace
109-
_configureServices?.Invoke(services);
110-
});
111-
112-
protected override void Dispose(bool disposing)
11360
{
114-
if (disposing)
61+
// Configure OpenTelemetry with in-memory exporters for all tests
62+
var otelBuilder = services.AddOpenTelemetry();
63+
_ = otelBuilder.WithTracing(tracing =>
11564
{
116-
foreach (var stream in _mockMemoryStreams)
117-
{
118-
stream.Dispose();
119-
}
120-
_mockMemoryStreams.Clear();
121-
}
122-
base.Dispose(disposing);
123-
}
65+
_ = tracing
66+
.AddDocsApiTracing() // Reuses production configuration
67+
.AddInMemoryExporter(ExportedActivities);
68+
});
69+
_ = otelBuilder.WithLogging(logging =>
70+
{
71+
_ = logging
72+
.AddDocsApiLogging() // Reuses production configuration
73+
.AddInMemoryExporter(ExportedLogRecords);
74+
});
75+
76+
// Mock IParameterProvider to avoid AWS dependencies in all tests
77+
var mockParameterProvider = A.Fake<IParameterProvider>();
78+
A.CallTo(() => mockParameterProvider.GetParam(A<string>._, A<bool>._, A<Cancel>._))
79+
.Returns(Task.FromResult("mock-value"));
80+
_ = services.AddSingleton(mockParameterProvider);
81+
82+
// Apply test-specific service replacements (if any)
83+
_configureServices?.Invoke(services);
84+
});
12485
}
12586

12687
/// <summary>
@@ -183,10 +144,10 @@ public ServiceReplacementBuilder ReplaceSingleton<TService>(TService instance) w
183144
/// Builds the final service configuration action.
184145
/// </summary>
185146
internal Action<IServiceCollection> Build() => services =>
186-
{
187-
foreach (var replacement in _replacements)
188-
{
189-
replacement(services);
190-
}
191-
};
147+
{
148+
foreach (var replacement in _replacements)
149+
{
150+
replacement(services);
151+
}
152+
};
192153
}

0 commit comments

Comments
 (0)