build-flow.api/tests/BuildFlow.Workflow.Tests/Reliability/ImportJobClaimTests.cs

using System.Collections.Concurrent;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging.Abstractions;
using Workflow.Application.BuildFlow.Imports.Commands;
using Workflow.Application.BuildFlow.Imports.Handlers;
using Workflow.Application.BuildFlow.Imports.Services;
using Workflow.Domain.Entities;
using Workflow.Domain.Enums;
using Xunit;

namespace BuildFlow.Workflow.Tests.Reliability;

/// <summary>
/// 导入任务的 PostgreSQL 原子领取和过期租约恢复测试。
/// </summary>
[Collection(PostgresWorkflowCollection.Name)]
public class ImportJobClaimTests(PostgresWorkflowFixture fixture)
{
    [Fact]
    public async Task TwoWorkersCanClaimPendingImportOnlyOnce()
    {
        await ResetAsync();
        var jobId = await SeedImportAsync(BuildFlowImportJobStatus.Pending);
        var processor = new RecordingImportJobProcessor();
        await using var dbA = fixture.CreateDbContext();
        await using var dbB = fixture.CreateDbContext();
        var handlerA = CreateHandler(dbA, processor);
        var handlerB = CreateHandler(dbB, processor);

        var results = await Task.WhenAll(
            handlerA.Handle(CreateCommand("worker-a"), CancellationToken.None),
            handlerB.Handle(CreateCommand("worker-b"), CancellationToken.None));

        results.Sum().Should().Be(1);
        processor.Calls.Should().ContainSingle(x => x.JobId == jobId);
        await using var verify = fixture.CreateDbContext();
        var job = await verify.BuildFlowImportJobs.SingleAsync(x => x.Id == jobId);
        job.Status.Should().Be(BuildFlowImportJobStatus.Running);
        job.AttemptCount.Should().Be(1);
        job.LeaseOwner.Should().BeOneOf("worker-a", "worker-b");
        job.LastHeartbeatAt.Should().NotBeNull();
    }

    [Fact]
    public async Task ExpiredRunningImportCanBeReclaimed()
    {
        await ResetAsync();
        var jobId = await SeedImportAsync(
            BuildFlowImportJobStatus.Running,
            "dead-worker",
            DateTime.UtcNow.AddMinutes(-1),
            attemptCount: 1);
        var processor = new RecordingImportJobProcessor();
        await using var db = fixture.CreateDbContext();
        var handler = CreateHandler(db, processor);

        var claimed = await handler.Handle(CreateCommand("recovery-worker"), CancellationToken.None);

        claimed.Should().Be(1);
        processor.Calls.Should().ContainSingle(x =>
            x.JobId == jobId && x.LeaseOwner == "recovery-worker");
        await using var verify = fixture.CreateDbContext();
        var job = await verify.BuildFlowImportJobs.SingleAsync(x => x.Id == jobId);
        job.AttemptCount.Should().Be(2);
        job.LeaseOwner.Should().Be("recovery-worker");
        job.LastHeartbeatAt.Should().NotBeNull();
    }

    [Fact]
    public async Task ExpiredImportAtMaxAttemptsBecomesFailed()
    {
        await ResetAsync();
        var jobId = await SeedImportAsync(
            BuildFlowImportJobStatus.Running,
            "dead-worker",
            DateTime.UtcNow.AddMinutes(-1),
            attemptCount: 5);
        var processor = new RecordingImportJobProcessor();
        await using var db = fixture.CreateDbContext();
        var handler = CreateHandler(db, processor);

        var claimed = await handler.Handle(CreateCommand("worker-a"), CancellationToken.None);

        claimed.Should().Be(0);
        processor.Calls.Should().BeEmpty();
        await using var verify = fixture.CreateDbContext();
        var job = await verify.BuildFlowImportJobs.SingleAsync(x => x.Id == jobId);
        job.Status.Should().Be(BuildFlowImportJobStatus.Failed);
        job.LeaseOwner.Should().BeNull();
        job.LeaseExpiresAt.Should().BeNull();
        job.FinishedAt.Should().NotBeNull();
        job.ErrorMessage.Should().Contain("最大重试次数");
    }

    private static ProcessAvailableImportJobsCommand CreateCommand(string owner) => new(
        owner,
        BatchSize: 10,
        LeaseDuration: TimeSpan.FromMinutes(5),
        MaxAttempts: 5);

    private static ProcessAvailableImportJobsCommandHandler CreateHandler(
        WorkflowDbContext db,
        IImportJobProcessor processor)
        => new(
            db,
            processor,
            NullLogger<ProcessAvailableImportJobsCommandHandler>.Instance);

    private async Task ResetAsync()
    {
        await using var db = fixture.CreateDbContext();
        await db.BuildFlowImportJobs.ExecuteDeleteAsync();
    }

    private async Task<Guid> SeedImportAsync(
        BuildFlowImportJobStatus status,
        string? leaseOwner = null,
        DateTime? leaseExpiresAt = null,
        int attemptCount = 0)
    {
        await using var db = fixture.CreateDbContext();
        var job = new BuildFlowImportJob
        {
            Id = Guid.NewGuid(),
            FormDefinitionId = Guid.NewGuid(),
            Status = status,
            FileUrl = "unused.xlsx",
            OriginalFileName = "unused.xlsx",
            LeaseOwner = leaseOwner,
            LeaseExpiresAt = leaseExpiresAt,
            AttemptCount = attemptCount,
        };
        db.BuildFlowImportJobs.Add(job);
        await db.SaveChangesAsync();
        return job.Id;
    }

    private sealed class RecordingImportJobProcessor : IImportJobProcessor
    {
        public ConcurrentQueue<ImportJobCall> Calls { get; } = new();

        public Task ProcessAsync(
            Guid jobId,
            string leaseOwner,
            TimeSpan leaseDuration,
            CancellationToken ct)
        {
            Calls.Enqueue(new ImportJobCall(jobId, leaseOwner, leaseDuration));
            return Task.CompletedTask;
        }
    }

    private sealed record ImportJobCall(
        Guid JobId,
        string LeaseOwner,
        TimeSpan LeaseDuration);
}