Register services in base class

test/IntegrationTests/NatsIntegrationFixture.cs

@@ -1,8 +1,9 @@
-using System;
-using System.Threading;
-using System.Threading.Tasks;
 using Aspire.Hosting;
+using CodeCargo.NatsDistributedCache.TestUtils;
+using Microsoft.Extensions.Logging;
 using NATS.Client.Core;
+using NATS.Client.KeyValueStore;
+using NATS.Net;
 
 namespace CodeCargo.NatsDistributedCache.IntegrationTests;
 
@@ -13,13 +14,14 @@ public class NatsIntegrationFixture : IAsyncLifetime
 {
     private static readonly TimeSpan StartupTimeout = TimeSpan.FromSeconds(30);
     private DistributedApplication? _app;
-    private INatsConnection? _natsConnection;
+    private int _disposed;
+    private ServiceProvider? _serviceProvider;
 
     /// <summary>
-    /// Gets the NATS connection for accessing the test NATS server
+    /// Gets the NATS connection
     /// </summary>
-    public INatsConnection NatsConnection => _natsConnection ?? throw new InvalidOperationException(
-        "NATS connection not initialized. Make sure InitializeAsync has been called.");
+    public INatsConnection NatsConnection => _serviceProvider?.GetRequiredService<INatsConnection>()
+                                             ?? throw new InvalidOperationException("InitializeAsync was not called");
 
     /// <summary>
     /// Initializes the fixture by starting the NATS server and creating a connection
@@ -43,31 +45,48 @@ public async ValueTask InitializeAsync()
             throw new InvalidOperationException("Cannot find connection string for NATS");
         }
 
-        // Create a NATS connection
-        var opts = new NatsOpts
+        // service provider
+        var services = new ServiceCollection();
+        services.AddLogging(builder =>
         {
-            Url = natsConnectionString,
-            Name = "IntegrationTest",
-            RequestReplyMode = NatsRequestReplyMode.Direct,
-        };
-        _natsConnection = new NatsConnection(opts);
+            builder.AddConsole();
+        });
+        services.AddNatsTestServices(natsConnectionString);
+        _serviceProvider = services.BuildServiceProvider();
+
+        // create the KV store
+        var kvContext = NatsConnection.CreateKeyValueStoreContext();
+        await kvContext.CreateOrUpdateStoreAsync(
+            new NatsKVConfig("cache") { LimitMarkerTTL = TimeSpan.FromSeconds(1), },
+            TestContext.Current.CancellationToken);
     }
 
+    /// <summary>
+    /// Configures the services with the NATS connection
+    /// </summary>
+    /// <param name="services">The service collection to configure</param>
+    public void ConfigureServices(IServiceCollection services) => services.AddSingleton(NatsConnection);
+
     /// <summary>
     /// Disposes the fixture by shutting down the NATS server
     /// </summary>
     public async ValueTask DisposeAsync()
     {
-        if (_natsConnection != null)
+        if (Interlocked.Increment(ref _disposed) != 1)
         {
-            await _natsConnection.DisposeAsync();
-            _natsConnection = null;
+            return;
         }
 
         if (_app != null)
         {
+            var stopCts = new CancellationTokenSource(TimeSpan.FromSeconds(5));
+            await _app.StopAsync(stopCts.Token);
             await _app.DisposeAsync();
-            _app = null;
+        }
+
+        if (_serviceProvider != null)
+        {
+            await _serviceProvider.DisposeAsync();
         }
 
         GC.SuppressFinalize(this);

test/IntegrationTests/TestBase.cs

@@ -1,55 +1,31 @@
-using System;
-using System.Collections.Concurrent;
-using System.Threading.Tasks;
 using CodeCargo.NatsDistributedCache.TestUtils.Services.Logging;
-using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
 using NATS.Client.Core;
-using NATS.Client.JetStream;
 using NATS.Client.KeyValueStore;
 using NATS.Net;
-using Xunit;
-using Xunit.Sdk;
 
 namespace CodeCargo.NatsDistributedCache.IntegrationTests;
 
 /// <summary>
 /// Base class for NATS integration tests that provides test output logging and fixture access
 /// </summary>
-public abstract class TestBase(NatsIntegrationFixture fixture) : IAsyncLifetime
+public abstract class TestBase : IAsyncDisposable
 {
-    private const string ServiceProviderKey = "ServiceProvider";
-    private readonly NatsIntegrationFixture _fixture = fixture ?? throw new ArgumentNullException(nameof(fixture));
-    private static readonly ConcurrentBag<ServiceProvider> _serviceProviders = new();
-    private ServiceProvider? _serviceProvider;
+    private int _disposed;
 
     /// <summary>
-    /// Gets the NATS connection from the fixture
+    /// Constructor that sets up the service provider with test output logging
     /// </summary>
-    protected INatsConnection NatsConnection => _fixture.NatsConnection;
-
-    /// <summary>
-    /// Gets the service provider configured with test logging and NATS services
-    /// </summary>
-    protected IServiceProvider ServiceProvider => _serviceProvider ??
-        throw new InvalidOperationException("Service provider not initialized. Make sure InitializeAsync has been called.");
-
-    // Static constructor to register for test class configuration
-    static TestBase()
-    {
-        // xUnit v3 will call ConfigureTestClass
-    }
-
-    /// <summary>
-    /// Configures the test class - this method is called by xUnit v3
-    /// </summary>
-    /// <param name="context">The test context</param>
-    public static void ConfigureTestClass(TestContext context)
+    /// <param name="fixture">The NATS integration fixture</param>
+    protected TestBase(NatsIntegrationFixture fixture)
     {
-        // Get the test output helper
-        var output = context.TestOutputHelper;
+        // Get the test output helper from the current test context
+        var testContext = TestContext.Current;
+        var output = testContext.TestOutputHelper;
         if (output == null)
-            return;
+        {
+            throw new InvalidOperationException("TestOutputHelper was not available in the current test context");
+        }
 
         // Create a service collection and configure logging
         var services = new ServiceCollection();
@@ -59,48 +35,34 @@ public static void ConfigureTestClass(TestContext context)
             builder.AddXUnitTestOutput(output);
         });
 
-        // Build the service provider and store it in the context's key-value storage
-        var serviceProvider = services.BuildServiceProvider();
-        _serviceProviders.Add(serviceProvider);
-        context.KeyValueStorage[ServiceProviderKey] = serviceProvider;
+        // Configure the service collection with NATS connection
+        fixture.ConfigureServices(services);
+
+        // Build service provider
+        ServiceProvider = services.BuildServiceProvider();
     }
 
     /// <summary>
-    /// Initializes the test by ensuring the KV store exists
+    /// Gets the service provider configured with test logging and NATS services
     /// </summary>
-    public virtual async ValueTask InitializeAsync()
-    {
-        // Setup service provider for this instance
-        var services = new ServiceCollection();
-
-        // Add NATS connection directly
-        services.AddSingleton<INatsConnection>(_fixture.NatsConnection);
-
-        // Add logging
-        services.AddLogging(builder =>
-        {
-            builder.ClearProviders();
-        });
+    protected ServiceProvider ServiceProvider { get; }
 
-        // Build service provider
-        _serviceProvider = services.BuildServiceProvider();
-        _serviceProviders.Add(_serviceProvider);
-
-        // Create or ensure KV store exists
-        var jsContext = NatsConnection.CreateJetStreamContext();
-        var kvContext = new NatsKVContext(jsContext);
-        await kvContext.CreateOrUpdateStoreAsync(new NatsKVConfig("cache"));
-    }
+    /// <summary>
+    /// Gets the NATS connection from the service provider
+    /// </summary>
+    protected INatsConnection NatsConnection => ServiceProvider.GetRequiredService<INatsConnection>();
 
     /// <summary>
     /// Cleanup after the test
     /// </summary>
     public virtual async ValueTask DisposeAsync()
     {
-        if (_serviceProvider != null)
+        if (Interlocked.Increment(ref _disposed) != 1)
         {
-            await _serviceProvider.DisposeAsync();
-            _serviceProvider = null;
+            return;
         }
+
+        await ServiceProvider.DisposeAsync();
+        GC.SuppressFinalize(this);
     }
 }

util/PerfTest/Program.cs

@@ -51,7 +51,9 @@
 Console.WriteLine("Creating KV store...");
 var nats = host.Services.GetRequiredService<INatsConnection>();
 var kv = nats.CreateKeyValueStoreContext();
-await kv.CreateStoreAsync(new NatsKVConfig("cache"), startupCts.Token);
+await kv.CreateOrUpdateStoreAsync(
+    new NatsKVConfig("cache") { LimitMarkerTTL = TimeSpan.FromSeconds(1) },
+    startupCts.Token);
 Console.WriteLine("KV store created");
 
 // Run the host