Adds GetUsageSnapshot/GetModelUsage/GetTaskUsage to WorkerHub (backed by a shared UsageSnapshotBuilder), a UsageUpdated broadcast fired after every UsageMonitorService poll cycle, and records the resolved model on each task_runs row so per-model/per-task usage can be reported from history.
104 lines
3.7 KiB
C#
104 lines
3.7 KiB
C#
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 static UsageSnapshot MakeSnapshot() => new(new UsageBucket(1, null), null, [], DateTime.UtcNow);
|
|
|
|
private (UsageMonitorService Service, UsageState State, CapturingHubContext Hub) CreateService(FakeClient client, WorkerConfig? cfg = 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);
|
|
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);
|
|
}
|
|
}
|