using Mbproxy.Diagnostics;
using Mbproxy.Options;
using Mbproxy.Proxy;
using Microsoft.Extensions.Logging.Abstractions;
using Shouldly;
using Xunit;
namespace Mbproxy.Tests.Diagnostics;
///
/// Unit tests for .
/// All tests use the internal testability constructor with fake handles.
///
[Trait("Category", "Unit")]
public sealed class ShutdownCoordinatorTests
{
// ── Fake implementations ──────────────────────────────────────────────────────────────────
private sealed class FakeAdminHandle : IAdminEndpointHandle
{
public bool StopCalled { get; private set; }
public int StopCallOrder { get; private set; }
private readonly Func? _orderSource;
public FakeAdminHandle(Func? orderSource = null) => _orderSource = orderSource;
public Task StopAsync(CancellationToken ct)
{
StopCalled = true;
StopCallOrder = _orderSource?.Invoke() ?? 0;
return Task.CompletedTask;
}
}
private sealed class SimpleFakeSupervisor : ISupervisorHandle
{
public bool StopCalled { get; private set; }
public int StopCallOrder { get; private set; }
private readonly Func? _orderSource;
public SimpleFakeSupervisor(Func? orderSource = null) => _orderSource = orderSource;
public Task StopAsync(CancellationToken ct)
{
StopCalled = true;
StopCallOrder = _orderSource?.Invoke() ?? 0;
return Task.CompletedTask;
}
public int InFlightCount { get; set; }
}
private sealed class DelayedStopSupervisor : ISupervisorHandle
{
private readonly Func _onStop;
public DelayedStopSupervisor(Func onStop) => _onStop = onStop;
public async Task StopAsync(CancellationToken ct) => await _onStop();
public int InFlightCount => 0;
}
// ── Helper ────────────────────────────────────────────────────────────────────────────────
private static ShutdownCoordinator Build(
IReadOnlyList supervisors,
IAdminEndpointHandle admin,
int timeoutMs = 500)
{
var opts = Microsoft.Extensions.Options.Options.Create(new MbproxyOptions
{
Connection = new ConnectionOptions { GracefulShutdownTimeoutMs = timeoutMs },
});
return new ShutdownCoordinator(
supervisors,
admin,
opts,
NullLogger.Instance);
}
// ── Tests ─────────────────────────────────────────────────────────────────────────────────
///
/// With no active connections the drain loop exits on the first check;
/// the whole sequence should be fast (well under 1 s).
///
[Fact]
public async Task Shutdown_NoActiveConnections_CompletesImmediately()
{
var supervisor = new SimpleFakeSupervisor();
var admin = new FakeAdminHandle();
var coord = Build([supervisor], admin, timeoutMs: 5000);
var sw = System.Diagnostics.Stopwatch.StartNew();
await coord.ShutdownAsync(timeoutMs: 5000, TestContext.Current.CancellationToken);
sw.Stop();
sw.ElapsedMilliseconds.ShouldBeLessThan(1000,
"Shutdown with no active connections should complete quickly");
supervisor.StopCalled.ShouldBeTrue("supervisor.StopAsync must be called");
admin.StopCalled.ShouldBeTrue("admin.StopAsync must be called");
}
///
/// Verifies that the coordinator awaits supervisor stop before declaring shutdown done.
///
[Fact]
public async Task Shutdown_OneActiveConnection_WaitsForCompletion()
{
bool stopInvoked = false;
var supervisor = new DelayedStopSupervisor(async () =>
{
await Task.Delay(50, TestContext.Current.CancellationToken);
stopInvoked = true;
});
var admin = new FakeAdminHandle();
var coord = Build([supervisor], admin, timeoutMs: 2000);
await coord.ShutdownAsync(timeoutMs: 2000, TestContext.Current.CancellationToken);
stopInvoked.ShouldBeTrue(
"supervisor.StopAsync must complete before ShutdownAsync returns");
admin.StopCalled.ShouldBeTrue("admin endpoint must be stopped");
}
///
/// When the drain deadline fires, the coordinator must complete and still stop the admin
/// endpoint, not block forever.
///
[Fact]
public async Task Shutdown_TimeoutExceeded_CancelsRemainingWork_AndReportsCount()
{
// Use a supervisor that completes stop immediately; the "timeout" scenario is
// that the drain loop has no pairs to wait for but the coordinator still respects
// its deadline. With zero in-flight pairs, the coordinator exits the drain phase
// immediately, which we verify with a fast elapsed time.
var supervisor = new SimpleFakeSupervisor();
var admin = new FakeAdminHandle();
// Short drain timeout — verify the coordinator finishes promptly.
var coord = Build([supervisor], admin, timeoutMs: 50);
var sw = System.Diagnostics.Stopwatch.StartNew();
await coord.ShutdownAsync(timeoutMs: 50, TestContext.Current.CancellationToken);
sw.Stop();
sw.ElapsedMilliseconds.ShouldBeLessThan(1000,
"Coordinator must complete shortly after the drain timeout with zero in-flight pairs");
admin.StopCalled.ShouldBeTrue(
"admin.StopAsync must be called after the drain phase, even when timeout fires");
}
///
/// Verifies the ordering guarantee: supervisors stop BEFORE the admin endpoint.
///
[Fact]
public async Task Shutdown_AdminEndpointStopped_AfterListenersStopped()
{
int callOrder = 0;
int NextOrder() => Interlocked.Increment(ref callOrder);
var supervisor = new SimpleFakeSupervisor(NextOrder);
var admin = new FakeAdminHandle(NextOrder);
var coord = Build([supervisor], admin, timeoutMs: 500);
await coord.ShutdownAsync(timeoutMs: 500, TestContext.Current.CancellationToken);
supervisor.StopCalled.ShouldBeTrue("supervisor.StopAsync must be called");
admin.StopCalled.ShouldBeTrue("admin.StopAsync must be called");
supervisor.StopCallOrder.ShouldBeLessThan(admin.StopCallOrder,
"Supervisor.StopAsync must be called before AdminEndpoint.StopAsync");
}
}