jobtrackingapp/JobTrackerApi.Tests/AiUsageMeterTests.cs

using JobTrackerApi.Data;
using JobTrackerApi.Models;
using JobTrackerApi.Services;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
using Moq;
using Xunit;

namespace JobTrackerApi.Tests;

public sealed class AiUsageMeterTests
{
    [Fact]
    public async Task Reservation_is_idempotent_and_success_can_replace_the_conservative_estimate()
    {
        await InFixtureAsync(async fixture =>
        {
            var entitlements = new AccountEntitlements(true, true, 1_000_000, 10, 10_000);
            var first = await fixture.Meter.ReserveAsync("owner", entitlements, "workspace", "same", "ai.workspace", 400, 100, default);
            var duplicate = await fixture.Meter.ReserveAsync("owner", entitlements, "workspace", "same", "ai.workspace", 800, 200, default);

            Assert.True(first.Created);
            Assert.False(duplicate.Created);
            Assert.Equal(first.Record.Id, duplicate.Record.Id);
            await fixture.Meter.FinalizeAsync(first.Record.Id, 80, 24, default);

            var usage = Assert.Single(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
            Assert.Equal(80, usage.InputCharacterCount);
            Assert.Equal(24, usage.OutputCharacterCount);
            Assert.Equal(26, usage.EstimatedTokenCount);
        });
    }

    [Fact]
    public async Task Monthly_call_and_token_limits_are_enforced_across_existing_sources()
    {
        await InFixtureAsync(async fixture =>
        {
            var tokenLimited = new AccountEntitlements(true, true, 1_000_000, 10, 500);
            await fixture.Meter.ReserveAsync("owner", tokenLimited, "operation", "one", "strategy.snapshot", 1_600, 400, default);
            var tokens = await Assert.ThrowsAsync<AiUsageLimitException>(() => fixture.Meter.ReserveAsync(
                "owner", tokenLimited, "workspace", "two", "ai.workspace", 800, 200, default));
            Assert.Equal("monthly_ai_tokens_exhausted", tokens.Code);

            var callLimited = new AccountEntitlements(true, true, 1_000_000, 1, 10_000);
            var calls = await Assert.ThrowsAsync<AiUsageLimitException>(() => fixture.Meter.ReserveAsync(
                "owner", callLimited, "workspace", "three", "ai.workspace", 40, 10, default));
            Assert.Equal("monthly_ai_calls_exhausted", calls.Code);
            Assert.Single(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
        });
    }

    [Fact]
    public async Task Totals_are_owner_scoped_and_survive_private_interaction_deletion()
    {
        await InFixtureAsync(async fixture =>
        {
            fixture.Db.AiUsageRecords.AddRange(
                Record("owner", "owner-source", 100),
                Record("other", "other-source", 900));
            var company = new Company { OwnerUserId = "owner", Name = "Usage test" };
            var application = new JobApplication
            {
                OwnerUserId = "owner", Company = company, JobTitle = "Usage test role",
            };
            fixture.Db.AiInteractions.Add(new AiInteraction
            {
                OwnerUserId = "owner", JobApplication = application, Module = "job-analysis",
                Title = "Private generation", Provider = "ollama",
                ResultJson = "{\"text\":\"private response\"}", CreatedAtUtc = DateTimeOffset.UtcNow,
            });
            await fixture.Db.SaveChangesAsync();

            var before = await fixture.Meter.AllTimeAsync("owner", default);
            Assert.Equal(new AiUsageTotals(1, 300, 100, 100), before);

            await fixture.Db.AiInteractions.ExecuteDeleteAsync();
            var after = await fixture.Meter.AllTimeAsync("owner", default);
            Assert.Equal(before, after);
        });
    }

    private static AiUsageRecord Record(string owner, string source, int tokens) => new()
    {
        OwnerUserId = owner,
        SourceType = "workspace",
        SourceId = source,
        TaskType = "ai.workspace",
        InputCharacterCount = tokens * 3,
        OutputCharacterCount = tokens,
        EstimatedTokenCount = tokens,
        CreatedAtUtc = DateTimeOffset.UtcNow,
    };

    private static async Task InFixtureAsync(Func<Fixture, Task> test)
    {
        await using var connection = new SqliteConnection("Data Source=:memory:");
        await connection.OpenAsync();
        var currentUser = new Mock<ICurrentUserService>();
        currentUser.SetupGet(item => item.UserId).Returns("owner");
        var options = new DbContextOptionsBuilder<JobTrackerContext>().UseSqlite(connection).Options;
        await using var db = new JobTrackerContext(options, currentUser.Object);
        await db.Database.EnsureCreatedAsync();
        await test(new Fixture(db, new AiUsageMeter(db, TimeProvider.System)));
    }

    private sealed record Fixture(JobTrackerContext Db, AiUsageMeter Meter);
}