jobtrackingapp/JobTrackerApi.Tests/AiOperationQueueTests.cs

using JobTrackerApi.Data;
using JobTrackerApi.Models;
using JobTrackerApi.Services;
using Microsoft.AspNetCore.Identity;
using Microsoft.Data.Sqlite;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.DependencyInjection;
using Xunit;

namespace JobTrackerApi.Tests;

public sealed class AiOperationQueueTests
{
    [Fact]
    public async Task Admission_is_pro_only_idempotent_bounded_and_contains_no_raw_payload()
    {
        await using var fixture = await Fixture.CreateAsync(perUserCapacity: 1);
        await fixture.SeedUserAsync("pro-1", pro: true);
        await using var scope = fixture.Provider.CreateAsyncScope();
        using var owner = scope.ServiceProvider.GetRequiredService<CurrentUserService>().UseBackgroundUser("pro-1");
        var admission = scope.ServiceProvider.GetRequiredService<AiOperationAdmission>();

        var first = await admission.EnqueueAsync("synthetic.ai", "same", "job", "42", AiOperationPriorities.Interactive, default);
        var duplicate = await admission.EnqueueAsync("synthetic.ai", "same", "job", "42", AiOperationPriorities.Interactive, default);
        var full = await Assert.ThrowsAsync<AiOperationAdmissionException>(() =>
            admission.EnqueueAsync("synthetic.ai", "other", "job", "43", AiOperationPriorities.Scheduled, default));

        Assert.True(first.Created);
        Assert.False(duplicate.Created);
        Assert.Equal(first.Operation.Id, duplicate.Operation.Id);
        Assert.Equal("/api/operations/" + first.Operation.Id.ToString("D"), first.StatusUrl);
        Assert.Equal("pro", first.Operation.EntitlementDecision);
        Assert.Equal("local_only", first.Operation.PrivacyPolicy);
        Assert.Equal("ai_queue_full", full.Code);
        Assert.Equal(15, full.RetryAfterSeconds);
        var usage = Assert.Single(await scope.ServiceProvider.GetRequiredService<JobTrackerContext>()
            .AiUsageRecords.IgnoreQueryFilters().AsNoTracking().ToListAsync());
        Assert.Equal(first.Operation.Id.ToString("D"), usage.SourceId);
        Assert.Equal(4_000, usage.EstimatedTokenCount);
    }

    [Fact]
    public async Task Free_and_ai_disabled_users_are_rejected_before_an_operation_is_created()
    {
        await using var fixture = await Fixture.CreateAsync();
        await fixture.SeedUserAsync("free-1", pro: false);
        await fixture.SeedUserAsync("disabled-1", pro: true, aiEnabled: false);

        Assert.Equal(ProEntitlement.RequiredCode, (await fixture.RejectedAsync("free-1")).Code);
        Assert.Equal(ProEntitlement.DisabledCode, (await fixture.RejectedAsync("disabled-1")).Code);
        await using var scope = fixture.Provider.CreateAsyncScope();
        Assert.Empty(await scope.ServiceProvider.GetRequiredService<JobTrackerContext>().UserOperations.IgnoreQueryFilters().ToListAsync());
    }

    [Fact]
    public async Task Worker_resolves_owner_completes_once_and_creates_persistent_notification()
    {
        await using var fixture = await Fixture.CreateAsync();
        await fixture.SeedUserAsync("pro-1", pro: true);
        await fixture.EnqueueAsync("pro-1", "run");

        Assert.True(await fixture.Provider.GetRequiredService<AiOperationWorker>().RunOnceAsync(default));
        Assert.False(await fixture.Provider.GetRequiredService<AiOperationWorker>().RunOnceAsync(default));

        await using var scope = fixture.Provider.CreateAsyncScope();
        var db = scope.ServiceProvider.GetRequiredService<JobTrackerContext>();
        var operation = await db.UserOperations.IgnoreQueryFilters().AsNoTracking().SingleAsync();
        Assert.Equal(OperationStatuses.Succeeded, operation.Status);
        Assert.Equal("synthetic-result", operation.ResultReference);
        Assert.Equal("ollama", operation.Provider);
        Assert.Equal("qwen-test", operation.Model);
        Assert.Equal("local_primary", operation.ProgressStage);
        Assert.Equal("pro-1", fixture.Handler.OwnerUserId);
        Assert.Equal("local_only", fixture.Handler.PrivacyPolicy);
        Assert.Equal("operation_succeeded", (await db.UserNotifications.IgnoreQueryFilters().AsNoTracking().SingleAsync()).Kind);
    }

    [Fact]
    public async Task Worker_classifies_retryable_failures_and_rechecks_live_entitlement()
    {
        await using var retryFixture = await Fixture.CreateAsync();
        await retryFixture.SeedUserAsync("pro-1", pro: true);
        await retryFixture.EnqueueAsync("pro-1", "retry");
        retryFixture.Handler.Failure = new AiOperationFailure("provider_unavailable", "Provider is unavailable.", retryable: true);
        Assert.True(await retryFixture.Provider.GetRequiredService<AiOperationWorker>().RunOnceAsync(default));
        await using (var scope = retryFixture.Provider.CreateAsyncScope())
        {
            var row = await scope.ServiceProvider.GetRequiredService<JobTrackerContext>().UserOperations.IgnoreQueryFilters().AsNoTracking().SingleAsync();
            Assert.Equal(OperationStatuses.WaitingForRetry, row.Status);
            Assert.Equal("provider_unavailable", row.FailureCategory);
        }

        await using var downgradeFixture = await Fixture.CreateAsync();
        await downgradeFixture.SeedUserAsync("pro-2", pro: true);
        await downgradeFixture.EnqueueAsync("pro-2", "downgrade");
        await downgradeFixture.SetAiEnabledAsync("pro-2", false);
        Assert.True(await downgradeFixture.Provider.GetRequiredService<AiOperationWorker>().RunOnceAsync(default));
        await using var verification = downgradeFixture.Provider.CreateAsyncScope();
        var failed = await verification.ServiceProvider.GetRequiredService<JobTrackerContext>().UserOperations.IgnoreQueryFilters().AsNoTracking().SingleAsync();
        Assert.Equal(OperationStatuses.Failed, failed.Status);
        Assert.Equal("entitlement_changed", failed.FailureCategory);
    }

    [Fact]
    public async Task Worker_records_provider_metadata_for_generation_failures()
    {
        await using var fixture = await Fixture.CreateAsync();
        await fixture.SeedUserAsync("pro-1", pro: true);
        await fixture.EnqueueAsync("pro-1", "provider-failure");
        fixture.Handler.GenerationFailure = new AiGenerationException(
            "provider_unavailable",
            "AI provider unavailable.",
            retryable: true,
            provider: "gemini",
            model: "gemini-test",
            routeReason: "external_fallback");

        Assert.True(await fixture.Provider.GetRequiredService<AiOperationWorker>().RunOnceAsync(default));

        await using var scope = fixture.Provider.CreateAsyncScope();
        var row = await scope.ServiceProvider.GetRequiredService<JobTrackerContext>()
            .UserOperations.IgnoreQueryFilters().AsNoTracking().SingleAsync();
        Assert.Equal(OperationStatuses.WaitingForRetry, row.Status);
        Assert.Equal("gemini", row.Provider);
        Assert.Equal("gemini-test", row.Model);
        Assert.Equal("external_fallback", row.ProgressStage);
    }

    private sealed class SyntheticHandler : IAiOperationHandler
    {
        public string TaskType => "synthetic.ai";
        public string? OwnerUserId { get; private set; }
        public string? PrivacyPolicy { get; private set; }
        public AiOperationFailure? Failure { get; set; }
        public AiGenerationException? GenerationFailure { get; set; }

        public Task<AiOperationExecutionResult> ExecuteAsync(AiOperationExecutionContext context, IServiceProvider services, CancellationToken cancellationToken)
        {
            if (Failure is not null) throw Failure;
            if (GenerationFailure is not null) throw GenerationFailure;
            OwnerUserId = services.GetRequiredService<ICurrentUserService>().UserId;
            PrivacyPolicy = context.EffectivePrivacyPolicy;
            return Task.FromResult(new AiOperationExecutionResult(
                "synthetic-result", "ollama", "qwen-test", "local_primary", 120, 40));
        }
    }

    private sealed class Fixture : IAsyncDisposable
    {
        private readonly SqliteConnection _connection;
        public ServiceProvider Provider { get; }
        public SyntheticHandler Handler { get; }

        private Fixture(SqliteConnection connection, ServiceProvider provider, SyntheticHandler handler)
        {
            _connection = connection;
            Provider = provider;
            Handler = handler;
        }

        public static async Task<Fixture> CreateAsync(int perUserCapacity = 10)
        {
            var connection = new SqliteConnection("Data Source=:memory:");
            await connection.OpenAsync();
            var configuration = new ConfigurationBuilder().AddInMemoryCollection(new Dictionary<string, string?>
            {
                ["AiQueue:PerUserCapacity"] = perUserCapacity.ToString(),
                ["AiQueue:GlobalCapacity"] = "100",
                ["AiQueue:HeartbeatSeconds"] = "5",
                ["Ai:ExternalProcessingEnabled"] = "false",
            }).Build();
            var handler = new SyntheticHandler();
            var services = new ServiceCollection();
            services.AddSingleton<IConfiguration>(configuration);
            services.AddLogging();
            services.AddHttpContextAccessor();
            services.AddScoped<CurrentUserService>();
            services.AddScoped<ICurrentUserService>(provider => provider.GetRequiredService<CurrentUserService>());
            services.AddDbContext<JobTrackerContext>((_, options) => options.UseSqlite(connection));
            services.AddIdentityCore<ApplicationUser>().AddRoles<IdentityRole>().AddEntityFrameworkStores<JobTrackerContext>();
            services.AddSingleton(TimeProvider.System);
            services.AddScoped<UserOperationStore>();
            services.AddSingleton<AiPrivacyPolicy>();
            services.AddSingleton<AiOperationExecutionScope>();
            services.AddScoped<AiUsageMeter>();
            services.AddScoped<AiOperationAdmission>();
            services.AddSingleton<IAiOperationHandler>(handler);
            services.AddSingleton<AiOperationWorker>();
            var provider = services.BuildServiceProvider();
            await using var scope = provider.CreateAsyncScope();
            await scope.ServiceProvider.GetRequiredService<JobTrackerContext>().Database.EnsureCreatedAsync();
            return new Fixture(connection, provider, handler);
        }

        public async Task SeedUserAsync(string userId, bool pro, bool aiEnabled = true)
        {
            await using var scope = Provider.CreateAsyncScope();
            var roles = scope.ServiceProvider.GetRequiredService<RoleManager<IdentityRole>>();
            if (pro && !await roles.RoleExistsAsync("Premium")) await roles.CreateAsync(new IdentityRole("Premium"));
            var users = scope.ServiceProvider.GetRequiredService<UserManager<ApplicationUser>>();
            var user = new ApplicationUser { Id = userId, UserName = $"{userId}@example.test", Email = $"{userId}@example.test", AiEnabled = aiEnabled };
            Assert.True((await users.CreateAsync(user)).Succeeded);
            if (pro) Assert.True((await users.AddToRoleAsync(user, "Premium")).Succeeded);
        }

        public async Task<AiOperationAdmissionException> RejectedAsync(string userId)
        {
            await using var scope = Provider.CreateAsyncScope();
            using var owner = scope.ServiceProvider.GetRequiredService<CurrentUserService>().UseBackgroundUser(userId);
            return await Assert.ThrowsAsync<AiOperationAdmissionException>(() => scope.ServiceProvider.GetRequiredService<AiOperationAdmission>()
                .EnqueueAsync("synthetic.ai", userId, "job", "42", AiOperationPriorities.Interactive, default));
        }

        public async Task EnqueueAsync(string userId, string key)
        {
            await using var scope = Provider.CreateAsyncScope();
            using var owner = scope.ServiceProvider.GetRequiredService<CurrentUserService>().UseBackgroundUser(userId);
            await scope.ServiceProvider.GetRequiredService<AiOperationAdmission>()
                .EnqueueAsync("synthetic.ai", key, "job", "42", AiOperationPriorities.Interactive, default);
        }

        public async Task SetAiEnabledAsync(string userId, bool enabled)
        {
            await using var scope = Provider.CreateAsyncScope();
            var users = scope.ServiceProvider.GetRequiredService<UserManager<ApplicationUser>>();
            var user = Assert.IsType<ApplicationUser>(await users.FindByIdAsync(userId));
            user.AiEnabled = enabled;
            Assert.True((await users.UpdateAsync(user)).Succeeded);
        }

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