ClaudeDo/tests/ClaudeDo.Worker.Tests/Usage/UsageMonitorServiceTests.cs

using ClaudeDo.Worker.Config;
using ClaudeDo.Worker.Hub;
using ClaudeDo.Worker.Tests.Infrastructure;
using ClaudeDo.Worker.Usage;
using ClaudeDo.Worker.Usage.Interfaces;
using Microsoft.Extensions.Logging.Abstractions;

namespace ClaudeDo.Worker.Tests.Usage;

public sealed class UsageMonitorServiceTests : IDisposable
{
    private readonly DbFixture _db = new();

    public void Dispose() => _db.Dispose();

    private sealed class FakeClient : IUsageClient
    {
        public Queue<Func<UsageSnapshot>> Results { get; } = new();
        public int CallCount { get; private set; }

        public Task<UsageSnapshot> GetUsageAsync(CancellationToken ct = default)
        {
            CallCount++;
            var next = Results.Count > 0 ? Results.Dequeue() : () => throw new InvalidOperationException("no result queued");
            return Task.FromResult(next());
        }
    }

    private sealed class FakeGate : IUsageGate
    {
        public Task<UsageGateDecision> EvaluateAsync(CancellationToken ct = default) =>
            Task.FromResult(new UsageGateDecision(false, null));
    }

    private sealed class FakeRunningProbe : IRunningTaskProbe
    {
        public bool AnyRunning { get; set; }
        public Task<bool> AnyRunningAsync(CancellationToken ct = default) => Task.FromResult(AnyRunning);
    }

    private static UsageSnapshot MakeSnapshot() => new(new UsageBucket(1, null), null, [], DateTime.UtcNow);

    private (UsageMonitorService Service, UsageState State, CapturingHubContext Hub) CreateService(
        FakeClient client, WorkerConfig? cfg = null, IRunningTaskProbe? probe = null)
    {
        var state = new UsageState();
        var config = cfg ?? new WorkerConfig();
        var builder = new UsageSnapshotBuilder(state, new FakeGate(), _db.CreateFactory(), config);
        var hubContext = new CapturingHubContext();
        var broadcaster = new HubBroadcaster(hubContext);
        var service = new UsageMonitorService(
            client, state, config, NullLogger<UsageMonitorService>.Instance, builder, broadcaster,
            probe ?? new FakeRunningProbe());
        return (service, state, hubContext);
    }

    [Fact]
    public async Task TickAsync_Success_UpdatesState()
    {
        var client = new FakeClient();
        client.Results.Enqueue(MakeSnapshot);
        var (service, state, _) = CreateService(client);

        await service.TickAsync(CancellationToken.None);

        Assert.NotNull(state.Snapshot);
        Assert.Null(state.LastError);
    }

    [Fact]
    public async Task TickAsync_Failure_SetsErrorButDoesNotThrow()
    {
        var client = new FakeClient();
        client.Results.Enqueue(() => throw new InvalidOperationException("network unreachable"));
        var (service, state, _) = CreateService(client);

        await service.TickAsync(CancellationToken.None);

        Assert.Equal("network unreachable", state.LastError);
        Assert.Null(state.Snapshot);
    }

    [Fact]
    public async Task TickAsync_FailureAfterSuccess_KeepsOldSnapshot()
    {
        var client = new FakeClient();
        client.Results.Enqueue(MakeSnapshot);
        client.Results.Enqueue(() => throw new InvalidOperationException("down"));
        var (service, state, _) = CreateService(client);

        await service.TickAsync(CancellationToken.None);
        await service.TickAsync(CancellationToken.None);

        Assert.NotNull(state.Snapshot);
        Assert.Equal("down", state.LastError);
    }

    [Fact]
    public async Task TickAsync_BroadcastsUsageUpdated_OnSuccessAndFailure()
    {
        var client = new FakeClient();
        client.Results.Enqueue(MakeSnapshot);
        client.Results.Enqueue(() => throw new InvalidOperationException("down"));
        var (service, _, hubContext) = CreateService(client);

        await service.TickAsync(CancellationToken.None);
        await service.TickAsync(CancellationToken.None);

        var calls = hubContext.Proxy.Calls.Where(c => c.Method == "UsageUpdated").ToList();
        Assert.Equal(2, calls.Count);
    }

    [Fact]
    public async Task RefreshNowAsync_Polls_And_Returns_Fresh_Snapshot()
    {
        var client = new FakeClient();
        client.Results.Enqueue(MakeSnapshot);
        var (service, _, _) = CreateService(client);

        var dto = await service.RefreshNowAsync();

        Assert.Equal(1, client.CallCount);
        Assert.False(dto.IsStale);
        Assert.NotNull(dto.FiveHourPercent);
    }

    [Fact]
    public async Task RefreshNowAsync_WithinCooldown_ReusesLastPoll()
    {
        var client = new FakeClient();
        client.Results.Enqueue(MakeSnapshot);
        var (service, _, _) = CreateService(client);

        await service.TickAsync(CancellationToken.None);
        var dto = await service.RefreshNowAsync();

        // Second call would have thrown "no result queued" had it hit the client.
        Assert.Equal(1, client.CallCount);
        Assert.False(dto.IsStale);
    }

    [Fact]
    public async Task RateLimit_Failure_Is_Recorded_As_Error()
    {
        var client = new FakeClient();
        client.Results.Enqueue(() => throw new UsageRateLimitedException(TimeSpan.FromSeconds(30)));
        var (service, state, _) = CreateService(client);

        await service.TickAsync(CancellationToken.None);

        Assert.NotNull(state.LastError);
        Assert.Contains("429", state.LastError);
    }
}

public sealed class UsagePollScheduleTests
{
    [Fact]
    public void Uses_active_interval_while_a_task_runs()
    {
        var delay = UsagePollSchedule.NextDelay(true, 300, 900, 0, null);
        Assert.Equal(TimeSpan.FromSeconds(300), delay);
    }

    [Fact]
    public void Uses_idle_interval_while_nothing_runs()
    {
        var delay = UsagePollSchedule.NextDelay(false, 300, 900, 0, null);
        Assert.Equal(TimeSpan.FromSeconds(900), delay);
    }

    [Fact]
    public void Backs_off_exponentially_on_repeated_rate_limits()
    {
        var first = UsagePollSchedule.NextDelay(true, 300, 900, 1, null);
        var second = UsagePollSchedule.NextDelay(true, 300, 900, 2, null);

        Assert.Equal(TimeSpan.FromSeconds(600), first);
        Assert.True(second > first);
    }

    [Fact]
    public void Honours_retry_after_when_longer_than_the_base_interval()
    {
        var delay = UsagePollSchedule.NextDelay(true, 300, 900, 1, TimeSpan.FromSeconds(420));
        Assert.Equal(TimeSpan.FromSeconds(420), delay);
    }

    [Fact]
    public void Never_polls_sooner_than_the_base_interval_after_a_rate_limit()
    {
        var delay = UsagePollSchedule.NextDelay(true, 300, 900, 1, TimeSpan.FromSeconds(5));
        Assert.Equal(TimeSpan.FromSeconds(300), delay);
    }

    [Fact]
    public void Caps_the_backoff()
    {
        var delay = UsagePollSchedule.NextDelay(false, 300, 900, 4, TimeSpan.FromHours(4));
        Assert.Equal(UsagePollSchedule.MaxDelay, delay);
    }
}