jobtrackingapp/JobTrackerApi.Tests/MeteredSummarizerServiceTests.cs

using System.Net;
using System.Text;
using JobTrackerApi.Data;
using JobTrackerApi.Models;
using JobTrackerApi.Services;
using JobTrackerApi.Tests.TestSupport;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Http;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Caching.Memory;
using Moq;
using Xunit;

namespace JobTrackerApi.Tests;

public sealed class MeteredSummarizerServiceTests
{
    [Fact]
    public async Task Usage_limit_handler_returns_stable_429_problem()
    {
        var context = new DefaultHttpContext { TraceIdentifier = "trace-ai-limit" };
        context.Response.Body = new MemoryStream();

        Assert.True(await new AiUsageLimitExceptionHandler().TryHandleAsync(
            context,
            new AiUsageLimitException("monthly_ai_calls_exhausted", "Monthly AI limit reached."),
            default));

        Assert.Equal(StatusCodes.Status429TooManyRequests, context.Response.StatusCode);
        context.Response.Body.Position = 0;
        using var reader = new StreamReader(context.Response.Body);
        var body = await reader.ReadToEndAsync();
        Assert.Contains("monthly_ai_calls_exhausted", body, StringComparison.Ordinal);
        Assert.Contains("trace-ai-limit", body, StringComparison.Ordinal);
    }

    [Fact]
    public async Task Successful_synchronous_generation_is_admitted_and_finalized()
    {
        await using var fixture = await Fixture.CreateAsync("{\"summary\":\"measured result\"}");

        var result = await fixture.Service.SummarizeAsync("measured input", 150, 30);

        Assert.Equal("measured result", result);
        var usage = Assert.Single(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
        Assert.Equal("synchronous", usage.SourceType);
        Assert.Equal("synchronous.summary", usage.TaskType);
        Assert.Equal("measured input".Length, usage.InputCharacterCount);
        Assert.Equal("measured result".Length, usage.OutputCharacterCount);
        Assert.Equal((usage.InputCharacterCount + usage.OutputCharacterCount + 3) / 4, usage.EstimatedTokenCount);
        Assert.Equal(1, fixture.Handler.RequestCount);
    }

    [Fact]
    public async Task Exhausted_limit_rejects_before_provider_call()
    {
        await using var fixture = await Fixture.CreateAsync("{\"summary\":\"must not run\"}");
        fixture.Db.AiUsageRecords.Add(new AiUsageRecord
        {
            OwnerUserId = "owner",
            SourceType = "synthetic",
            SourceId = "monthly-limit",
            TaskType = "synthetic",
            CallCount = 250,
            EstimatedTokenCount = 1,
            CreatedAtUtc = DateTimeOffset.UtcNow,
        });
        await fixture.Db.SaveChangesAsync();

        var error = await Assert.ThrowsAsync<AiUsageLimitException>(() => fixture.Service.SummarizeAsync("blocked"));

        Assert.Equal("monthly_ai_calls_exhausted", error.Code);
        Assert.Equal(0, fixture.Handler.RequestCount);
    }

    [Fact]
    public async Task Existing_metered_scope_bypasses_the_synchronous_decorator()
    {
        await using var fixture = await Fixture.CreateAsync("{\"summary\":\"already metered\"}");

        using (fixture.UsageScope.Suppress())
            Assert.Equal("already metered", await fixture.Service.SummarizeAsync("workspace input"));

        Assert.Empty(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
        Assert.Equal(1, fixture.Handler.RequestCount);
    }

    [Fact]
    public async Task Durable_operation_scope_does_not_create_a_second_usage_record()
    {
        await using var fixture = await Fixture.CreateAsync("{\"summary\":\"operation result\"}");
        var lease = new UserOperationLease(
            Guid.NewGuid(), "owner", "lease", "strategy.snapshot", "local_only", "job", "1", 1, null);

        using (fixture.OperationScope.Use(new AiOperationExecutionContext(lease, "local_only")))
            Assert.Equal("operation result", await fixture.Service.SummarizeAsync("operation input"));

        Assert.Empty(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
        Assert.Equal(1, fixture.Handler.RequestCount);
    }

    [Fact]
    public async Task Free_owner_is_rejected_at_the_shared_provider_boundary()
    {
        await using var fixture = await Fixture.CreateAsync("{\"summary\":\"must not run\"}", Array.Empty<string>());

        var error = await Assert.ThrowsAsync<AiUsageLimitException>(() => fixture.Service.SummarizeAsync("blocked"));

        Assert.Equal("ai_not_available", error.Code);
        Assert.Empty(await fixture.Db.AiUsageRecords.AsNoTracking().ToListAsync());
        Assert.Equal(0, fixture.Handler.RequestCount);
    }

    private sealed class Fixture : IAsyncDisposable
    {
        private readonly SqliteConnection _connection;
        public JobTrackerContext Db { get; }
        public MeteredSummarizerService Service { get; }
        public CountingHandler Handler { get; }
        public AiUsageExecutionScope UsageScope { get; }
        public AiOperationExecutionScope OperationScope { get; }

        private Fixture(
            SqliteConnection connection,
            JobTrackerContext db,
            MeteredSummarizerService service,
            CountingHandler handler,
            AiUsageExecutionScope usageScope,
            AiOperationExecutionScope operationScope)
        {
            _connection = connection;
            Db = db;
            Service = service;
            Handler = handler;
            UsageScope = usageScope;
            OperationScope = operationScope;
        }

        public static async Task<Fixture> CreateAsync(string responseJson, IReadOnlyList<string>? roles = null)
        {
            var connection = new SqliteConnection("Data Source=:memory:");
            await connection.OpenAsync();
            var currentUser = new Mock<ICurrentUserService>();
            currentUser.SetupGet(service => service.UserId).Returns("owner");
            var db = new JobTrackerContext(
                new DbContextOptionsBuilder<JobTrackerContext>().UseSqlite(connection).Options,
                currentUser.Object);
            await db.Database.EnsureCreatedAsync();

            var handler = new CountingHandler(responseJson);
            var client = new HttpClient(handler) { BaseAddress = new Uri("http://ai.test") };
            var factory = new Mock<IHttpClientFactory>();
            factory.Setup(item => item.CreateClient("ai-service")).Returns(client);
            var inner = new SummarizerService(factory.Object, new MemoryCache(new MemoryCacheOptions()));
            var user = new ApplicationUser { Id = "owner", AiEnabled = true };
            var users = TestHostFactory.CreateUserManager(user);
            users.Setup(item => item.GetRolesAsync(user)).ReturnsAsync((roles ?? new[] { "Premium" }).ToList());
            var usageScope = new AiUsageExecutionScope();
            var operationScope = new AiOperationExecutionScope();
            var service = new MeteredSummarizerService(
                inner,
                db,
                users.Object,
                new AiUsageMeter(db, TimeProvider.System),
                operationScope,
                usageScope);
            return new Fixture(connection, db, service, handler, usageScope, operationScope);
        }

        public async ValueTask DisposeAsync()
        {
            await Db.DisposeAsync();
            await _connection.DisposeAsync();
        }
    }

    public sealed class CountingHandler(string responseJson) : HttpMessageHandler
    {
        public int RequestCount { get; private set; }

        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            RequestCount++;
            return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
            {
                Content = new StringContent(responseJson, Encoding.UTF8, "application/json"),
            });
        }
    }
}