w4c-workflows-api/w4c-workflows-api.Tests/RunLifecycleEngineTests.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Logging.Abstractions;
using w4c_workflows.Data;
using w4c_workflows.Models;
using w4c_workflows.Services.Execution;
using w4c_workflows.Services.Runs;
using Xunit;

namespace w4c_workflows.Tests;

[Collection("WorkflowsPostgres")]
public class RunLifecycleEngineTests
{
    private readonly WorkflowsPostgresFixture _fixture;

    public RunLifecycleEngineTests(WorkflowsPostgresFixture fixture)
    {
        _fixture = fixture;
    }

    private static RunLifecycleEngine NewEngine(FakeEventBus events, FakeLeaseService leases, int maxRetries = 3, int retryBaseDelayMs = 0)
        => new(events, leases,
            new ConfigurationBuilder()
                .AddInMemoryCollection(new Dictionary<string, string?>
                {
                    ["Workflows:MaxRetries"] = maxRetries.ToString(),
                    ["Workflows:RetryBaseDelayMs"] = retryBaseDelayMs.ToString(),
                })
                .Build(),
            NullLogger<RunLifecycleEngine>.Instance);

    private static TaskDispatcher NewDispatcher(WorkflowsDbContext db, FakeJobQueue jobs)
        => new(db, jobs, new ConfigurationBuilder().Build(), NullLogger<TaskDispatcher>.Instance);

    private static IReadOnlyDictionary<string, string> ResultFields(
        WorkflowRun run, WorkflowTask task, int attempt, bool success, string? output = null, string? error = null)
    {
        var invocation = new TaskInvocation(
            TaskInvocation.TypeValue, run.Id.ToString(), task.Id.ToString(), task.Key,
            task.Language, "x", null, new Dictionary<string, string>(), null, null, run.TenantId, attempt);
        var result = new ExecutionResult(success, success ? 0 : 1, string.Empty, string.Empty, output, error, TimeSpan.Zero);
        return TaskResultMessage.ToFields(invocation, result);
    }

    /// <summary>The engine is multi-tenant: it scans the whole DB. The shared Postgres
    /// fixture may contain other tests' tenants, so assert on this tenant's jobs only.</summary>
    private static List<IReadOnlyDictionary<string, string>> EnqueuedFor(FakeJobQueue jobs, string tenantId)
        => jobs.Enqueued.Where(x => x.TenantId == tenantId).Select(x => x.Fields).ToList();

    private static string Tenant() => "t" + Guid.NewGuid().ToString("N")[..12];

    [Fact]
    public async Task Dispatch_marks_run_running_and_enqueues_root()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow, "{\"in\":1}");

        var jobs = new FakeJobQueue();
        var engine = NewEngine(new FakeEventBus(), new FakeLeaseService());

        var dispatched = await engine.DispatchPendingAsync(db, NewDispatcher(db, jobs), default);

        Assert.True(dispatched >= 1);

        var reloaded = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Running, reloaded.Status);
        Assert.NotNull(reloaded.StartedAt);

        var fields = Assert.Single(EnqueuedFor(jobs, tenantId));
        Assert.Equal("root", fields["task_key"]);
        Assert.Equal("{\"in\":1}", fields["input"]);
        Assert.Equal("1", fields["attempt"]);
    }

    [Fact]
    public async Task Full_chain_runs_to_completion()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService());
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");
        var b = WorkflowDataHelpers.TaskByKey(workflow, "b");

        events.SeedResult(tenantId, ResultFields(run, root, 1, true, "{\"out\":1}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var mine = EnqueuedFor(jobs, tenantId);
        Assert.Equal(2, mine.Count);
        Assert.Equal("a", mine[1]["task_key"]);
        Assert.Equal("{\"out\":1}", mine[1]["input"]);
        Assert.Equal(RunStatus.Running, (await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id)).Status);

        events.SeedResult(tenantId, ResultFields(run, a, 1, true, "{\"out\":2}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        mine = EnqueuedFor(jobs, tenantId);
        Assert.Equal(3, mine.Count);
        Assert.Equal("b", mine[2]["task_key"]);
        Assert.Equal("{\"out\":2}", mine[2]["input"]);

        events.SeedResult(tenantId, ResultFields(run, b, 1, true, "{\"out\":3}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var final = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Succeeded, final.Status);
        Assert.Equal("{\"out\":3}", final.OutputJson);
        Assert.NotNull(final.FinishedAt);
        Assert.Equal(3, EnqueuedFor(jobs, tenantId).Count); // no further dispatch

        var taskRuns = await db.TaskRuns.Where(t => t.RunId == run.Id).ToListAsync();
        Assert.Equal(3, taskRuns.Count);
        Assert.All(taskRuns, t => Assert.Equal(TaskRunStatus.Succeeded, t.Status));
    }

    [Fact]
    public async Task Failed_task_without_compensation_is_dead_lettered()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService(), maxRetries: 0);
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        events.SeedResult(tenantId, ResultFields(run, root, 1, false, error: "boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var final = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Dead, final.Status);
        Assert.Equal("boom", final.Error);
        Assert.NotNull(final.FinishedAt);

        var rootRun = await db.TaskRuns.SingleAsync(t => t.RunId == run.Id && t.TaskId == root.Id);
        Assert.Equal(TaskRunStatus.Dead, rootRun.Status);

        var dlq = Assert.Single(jobs.DeadLetteredFor(tenantId));
        Assert.Equal("max_retries_exceeded", dlq.Reason);
        Assert.Equal("boom", dlq.Fields["dlq.error"]);
        Assert.Equal("root", dlq.Fields["task_key"]);
    }

    [Fact]
    public async Task Failed_task_retries_with_backoff_then_dead_letters()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService(), maxRetries: 2);
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");

        // Attempt 1 fails → retry scheduled; run stays Running.
        events.SeedResult(tenantId, ResultFields(run, root, 1, false, error: "boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);
        Assert.Equal(RunStatus.Running, (await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id)).Status);
        Assert.NotNull((await db.TaskRuns.SingleAsync(t => t.RunId == run.Id && t.Attempt == 1)).NextAttemptAt);

        // Retry fires attempt 2.
        await engine.RetryDueAsync(db, dispatcher, default);
        Assert.Equal(2, jobs.EnqueuedFor(tenantId).Count);

        // Attempt 2 fails → retry scheduled again.
        events.SeedResult(tenantId, ResultFields(run, root, 2, false, error: "boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        // Retry fires attempt 3 (maxRetries=2 → 3 total attempts).
        await engine.RetryDueAsync(db, dispatcher, default);
        Assert.Equal(3, jobs.EnqueuedFor(tenantId).Count);

        // Attempt 3 fails → retries exhausted → dead-letter.
        events.SeedResult(tenantId, ResultFields(run, root, 3, false, error: "boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        Assert.Equal(RunStatus.Dead, (await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id)).Status);
        Assert.Single(jobs.DeadLetteredFor(tenantId));

        var attempts = await db.TaskRuns.Where(t => t.RunId == run.Id).OrderBy(t => t.Attempt).ToListAsync();
        Assert.Equal(3, attempts.Count);
    }

    [Fact]
    public async Task Failed_task_with_onError_runs_compensation_and_marks_run_failed()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.SagaYaml, "workflows/saga.yaml");
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService(), maxRetries: 0);
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");
        var compA = WorkflowDataHelpers.TaskByKey(workflow, "compensate_a");

        // root succeeds → dispatches a.
        events.SeedResult(tenantId, ResultFields(run, root, 1, true, "{\"out\":1}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        // a fails → retries exhausted → run compensates via a.onError.
        events.SeedResult(tenantId, ResultFields(run, a, 1, false, error: "a boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        Assert.Equal(RunStatus.Compensating, (await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id)).Status);
        Assert.Equal("compensate_a", jobs.EnqueuedFor(tenantId)[^1]["task_key"]);

        // Compensation succeeds → run Failed (compensated).
        events.SeedResult(tenantId, ResultFields(run, compA, 1, true, "{}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var final = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Failed, final.Status);
        Assert.Equal("a boom", final.Error);

        var compRun = await db.TaskRuns.SingleAsync(t => t.RunId == run.Id && t.TaskId == compA.Id);
        Assert.Equal(TaskRunStatus.Compensated, compRun.Status);
        Assert.True(compRun.IsCompensation);

        Assert.Empty(jobs.DeadLetteredFor(tenantId));
    }

    [Fact]
    public async Task Reverse_compensation_walks_succeeded_prefix_in_reverse()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.SagaYaml, "workflows/saga.yaml");
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService(), maxRetries: 0);
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");
        var b = WorkflowDataHelpers.TaskByKey(workflow, "b");
        var c = WorkflowDataHelpers.TaskByKey(workflow, "c");

        // Drive the success chain root → a → b → c.
        events.SeedResult(tenantId, ResultFields(run, root, 1, true, "{\"out\":1}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);
        events.SeedResult(tenantId, ResultFields(run, a, 1, true, "{\"out\":2}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);
        events.SeedResult(tenantId, ResultFields(run, b, 1, true, "{\"out\":3}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        // c fails → plan = [compensate_c (own onError), compensate_b, compensate_a].
        events.SeedResult(tenantId, ResultFields(run, c, 1, false, error: "c boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var compC = WorkflowDataHelpers.TaskByKey(workflow, "compensate_c");
        var compB = WorkflowDataHelpers.TaskByKey(workflow, "compensate_b");
        var compA = WorkflowDataHelpers.TaskByKey(workflow, "compensate_a");

        Assert.Equal("compensate_c", jobs.EnqueuedFor(tenantId)[^1]["task_key"]);

        events.SeedResult(tenantId, ResultFields(run, compC, 1, true, "{}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);
        Assert.Equal("compensate_b", jobs.EnqueuedFor(tenantId)[^1]["task_key"]);

        events.SeedResult(tenantId, ResultFields(run, compB, 1, true, "{}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);
        Assert.Equal("compensate_a", jobs.EnqueuedFor(tenantId)[^1]["task_key"]);

        events.SeedResult(tenantId, ResultFields(run, compA, 1, true, "{}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        Assert.Equal(RunStatus.Failed, (await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id)).Status);
        Assert.Empty(jobs.DeadLetteredFor(tenantId));
    }

    [Fact]
    public async Task Failed_compensation_dead_letters_the_run()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.SagaYaml, "workflows/saga.yaml");
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService(), maxRetries: 0);
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");
        var compA = WorkflowDataHelpers.TaskByKey(workflow, "compensate_a");

        events.SeedResult(tenantId, ResultFields(run, root, 1, true, "{\"out\":1}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        events.SeedResult(tenantId, ResultFields(run, a, 1, false, error: "a boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        // The compensation step itself fails → dead-letter.
        events.SeedResult(tenantId, ResultFields(run, compA, 1, false, error: "compensation boom"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var final = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Dead, final.Status);
        Assert.Equal("compensation boom", final.Error);

        var compRun = await db.TaskRuns.SingleAsync(t => t.RunId == run.Id && t.TaskId == compA.Id);
        Assert.Equal(TaskRunStatus.Dead, compRun.Status);

        var dlq = Assert.Single(jobs.DeadLetteredFor(tenantId));
        Assert.Equal("max_retries_exceeded", dlq.Reason);
        Assert.Equal("compensation boom", dlq.Fields["dlq.error"]);
        Assert.Equal("compensate_a", dlq.Fields["task_key"]);
    }

    [Fact]
    public async Task Durable_run_checkpoints_then_clears_on_completion()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.DurableYaml, "workflows/durable.yaml");
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService());
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");

        events.SeedResult(tenantId, ResultFields(run, root, 1, true, "{\"s\":1}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        // Checkpoint recorded: instance = run id, next task = a, input = root output.
        var checkpoint = await db.DurableStates.SingleAsync(d => d.InstanceId == run.Id.ToString());
        Assert.Equal(a.Id, checkpoint.TaskId);
        Assert.Equal("{\"s\":1}", checkpoint.StateJson);
        Assert.Equal("test-correlation", checkpoint.CorrelationId);

        events.SeedResult(tenantId, ResultFields(run, a, 1, true, "{\"s\":2}"));
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        var final = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Succeeded, final.Status);
        Assert.False(await db.DurableStates.AnyAsync(d => d.InstanceId == run.Id.ToString()));
    }

    [Fact]
    public async Task Resume_dispatches_the_start_task_instead_of_root()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var a = WorkflowDataHelpers.TaskByKey(workflow, "a");
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow, "{\"resume\":true}", a.Id);

        var jobs = new FakeJobQueue();
        var engine = NewEngine(new FakeEventBus(), new FakeLeaseService());

        await engine.DispatchPendingAsync(db, NewDispatcher(db, jobs), default);

        var fields = Assert.Single(EnqueuedFor(jobs, tenantId));
        Assert.Equal("a", fields["task_key"]);
        Assert.Equal("{\"resume\":true}", fields["input"]);

        var reloaded = await db.WorkflowRuns.SingleAsync(r => r.Id == run.Id);
        Assert.Equal(RunStatus.Running, reloaded.Status);
    }

    [Fact]
    public async Task Duplicate_result_is_ignored()
    {
        var tenantId = Tenant();
        await using var db = _fixture.CreateContext();
        var workflow = WorkflowDataHelpers.CompileAndSave(db, tenantId, WorkflowDataHelpers.LinearYaml);
        var run = WorkflowDataHelpers.AddPendingRun(db, workflow);

        var jobs = new FakeJobQueue();
        var events = new FakeEventBus();
        var engine = NewEngine(events, new FakeLeaseService());
        var dispatcher = NewDispatcher(db, jobs);
        var store = new DurableStateStore(db);

        await engine.DispatchPendingAsync(db, dispatcher, default);

        var root = WorkflowDataHelpers.TaskByKey(workflow, "root");
        var rootResult = ResultFields(run, root, 1, true, "{\"out\":1}");

        events.SeedResult(tenantId, rootResult);
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        Assert.Equal(2, EnqueuedFor(jobs, tenantId).Count); // root + a

        // Replay the exact same root result — must be a no-op.
        events.SeedResult(tenantId, rootResult);
        await engine.ProcessResultsAsync(db, dispatcher, store, default);

        Assert.Equal(2, EnqueuedFor(jobs, tenantId).Count); // no new dispatch
        var rootRun = await db.TaskRuns.SingleAsync(t => t.TaskId == root.Id);
        Assert.Equal(TaskRunStatus.Succeeded, rootRun.Status);
    }
}