ClaudeDo/tests/ClaudeDo.Worker.Tests/State/TaskStateServiceTests.cs

using ClaudeDo.Data;
using ClaudeDo.Data.Models;
using ClaudeDo.Data.Repositories;
using ClaudeDo.Worker.Planning;
using ClaudeDo.Worker.State;
using ClaudeDo.Worker.Tests.Infrastructure;
using Microsoft.EntityFrameworkCore;
using TaskStatus = ClaudeDo.Data.Models.TaskStatus;

namespace ClaudeDo.Worker.Tests.State;

file sealed class FakeActiveMergeState : IActiveMergeState
{
    public HashSet<string> ActiveTaskIds { get; } = new();
    public bool HasActiveMerge(string taskId) => ActiveTaskIds.Contains(taskId);
}

public sealed class TaskStateServiceTests : IDisposable
{
    private readonly DbFixture _db = new();
    private readonly TestDbContextFactory _factory;
    private readonly TaskStateServiceBuilder.Built _built;
    private readonly ITaskStateService _sut;
    private readonly string _listId;

    public TaskStateServiceTests()
    {
        _factory = _db.CreateFactory();
        _built = TaskStateServiceBuilder.Build(_factory);
        _sut = _built.State;

        _listId = Guid.NewGuid().ToString();
        using var ctx = _factory.CreateDbContext();
        ctx.Lists.Add(new ListEntity
        {
            Id = _listId,
            Name = "Test",
            CreatedAt = DateTime.UtcNow,
            DefaultCommitType = "chore",
        });
        ctx.SaveChanges();
    }

    public void Dispose() => _db.Dispose();

    private async Task<string> SeedTaskAsync(
        TaskStatus status,
        string? parentId = null,
        int sortOrder = 0,
        string? blockedBy = null,
        PlanningPhase phase = PlanningPhase.None,
        bool isManual = false)
    {
        var id = Guid.NewGuid().ToString();
        await using var ctx = _factory.CreateDbContext();
        ctx.Tasks.Add(new TaskEntity
        {
            Id = id,
            ListId = _listId,
            Title = "task",
            Status = status,
            PlanningPhase = phase,
            CreatedAt = DateTime.UtcNow,
            ParentTaskId = parentId,
            SortOrder = sortOrder,
            BlockedByTaskId = blockedBy,
            IsManual = isManual,
        });
        await ctx.SaveChangesAsync();
        return id;
    }

    private async Task<TaskStatus> GetStatusAsync(string id)
    {
        await using var ctx = _factory.CreateDbContext();
        return await ctx.Tasks.Where(t => t.Id == id).Select(t => t.Status).FirstAsync();
    }

    private async Task<TaskEntity> GetTaskAsync(string id)
    {
        await using var ctx = _factory.CreateDbContext();
        return await new TaskRepository(ctx).GetByIdAsync(id) ?? throw new InvalidOperationException($"task {id} not found");
    }

    // ─── EnqueueAsync ─────────────────────────────────────────────────────

    [Fact]
    public async Task EnqueueAsync_FromIdle_TransitionsToQueued_AndWakesQueue()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle);
        var wakesBefore = _built.WakeCount();

        var result = await _sut.EnqueueAsync(id, default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Queued, await GetStatusAsync(id));
        Assert.True(_built.WakeCount() > wakesBefore);
        Assert.Contains(_built.Hub.Proxy.Calls, c => c.Method == "TaskUpdated");
    }

    [Fact]
    public async Task EnqueueAsync_ManualTask_Rejected_AndStaysIdle()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle, isManual: true);

        var result = await _sut.EnqueueAsync(id, default);

        Assert.False(result.Ok);
        Assert.Contains("Manual task", result.Reason);
        Assert.Equal(TaskStatus.Idle, await GetStatusAsync(id));
    }

    [Fact]
    public async Task StartRunningAsync_ManualTask_Rejected_AndStaysIdle()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle, isManual: true);

        var result = await _sut.StartRunningAsync(id, DateTime.UtcNow, default);

        Assert.False(result.Ok);
        Assert.Contains("Manual task", result.Reason);
        Assert.Equal(TaskStatus.Idle, await GetStatusAsync(id));
    }

    [Fact]
    public async Task EnqueueAsync_FromRunning_Rejects_AndDoesNotMutate()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.EnqueueAsync(id, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Running, await GetStatusAsync(id));
    }

    [Fact]
    public async Task EnqueueAsync_DraftChild_Rejected_WhenParentNotFinalized()
    {
        var parent = await SeedTaskAsync(TaskStatus.Idle, phase: PlanningPhase.Active);
        var child = await SeedTaskAsync(TaskStatus.Idle, parentId: parent);

        var result = await _sut.EnqueueAsync(child, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Idle, await GetStatusAsync(child));
    }

    [Fact]
    public async Task EnqueueAsync_PlannedChild_Succeeds_WhenParentFinalized()
    {
        var parent = await SeedTaskAsync(TaskStatus.Idle, phase: PlanningPhase.Finalized);
        var child = await SeedTaskAsync(TaskStatus.Idle, parentId: parent);

        var result = await _sut.EnqueueAsync(child, default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Queued, await GetStatusAsync(child));
    }

    // ─── StartRunningAsync ────────────────────────────────────────────────

    [Fact]
    public async Task StartRunningAsync_FromQueued_TransitionsToRunning_AndStampsStartedAt()
    {
        var id = await SeedTaskAsync(TaskStatus.Queued);
        var startedAt = new DateTime(2026, 4, 27, 10, 0, 0, DateTimeKind.Utc);

        var result = await _sut.StartRunningAsync(id, startedAt, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Running, t.Status);
        Assert.Equal(startedAt, t.StartedAt);
    }

    [Fact]
    public async Task StartRunningAsync_FromRunning_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.StartRunningAsync(id, DateTime.UtcNow, default);

        Assert.False(result.Ok);
    }

    [Fact]
    public async Task StartRunningAsync_TwoParallelClaims_ExactlyOneWins()
    {
        var id = await SeedTaskAsync(TaskStatus.Queued);
        var startedAt = DateTime.UtcNow;

        // Two concurrent calls: only one ExecuteUpdate should affect a row.
        var t1 = Task.Run(() => _sut.StartRunningAsync(id, startedAt, default));
        var t2 = Task.Run(() => _sut.StartRunningAsync(id, startedAt, default));
        var results = await Task.WhenAll(t1, t2);

        var winners = results.Count(r => r.Ok);
        Assert.Equal(1, winners);
        Assert.Equal(TaskStatus.Running, await GetStatusAsync(id));
    }

    [Fact]
    public async Task StartRunningAsync_DraftChild_Rejected_WhenParentNotFinalized()
    {
        var parent = await SeedTaskAsync(TaskStatus.Idle, phase: PlanningPhase.Active);
        var child = await SeedTaskAsync(TaskStatus.Idle, parentId: parent);

        var result = await _sut.StartRunningAsync(child, DateTime.UtcNow, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Idle, await GetStatusAsync(child));
    }

    [Fact]
    public async Task StartRunningAsync_PlannedChild_Succeeds_WhenParentFinalized()
    {
        var parent = await SeedTaskAsync(TaskStatus.Idle, phase: PlanningPhase.Finalized);
        var child = await SeedTaskAsync(TaskStatus.Idle, parentId: parent);

        var result = await _sut.StartRunningAsync(child, DateTime.UtcNow, default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Running, await GetStatusAsync(child));
    }

    // ─── CompleteAsync ────────────────────────────────────────────────────

    [Fact]
    public async Task CompleteAsync_FromRunning_TransitionsToDone()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.CompleteAsync(id, DateTime.UtcNow, "ok", default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Done, t.Status);
        Assert.Equal("ok", t.Result);
        Assert.NotNull(t.FinishedAt);
    }

    [Fact]
    public async Task CompleteAsync_FromQueued_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Queued);

        var result = await _sut.CompleteAsync(id, DateTime.UtcNow, "ok", default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Queued, await GetStatusAsync(id));
    }

    // ─── FailAsync ────────────────────────────────────────────────────────

    [Fact]
    public async Task FailAsync_FromRunning_TransitionsToFailed()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.FailAsync(id, DateTime.UtcNow, "boom", default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Failed, t.Status);
        Assert.Equal("boom", t.Result);
    }

    [Fact]
    public async Task FailAsync_FromQueued_TransitionsToFailed()
    {
        // OverrideSlotService can call MarkFailed before StartRunningAsync when a
        // preflight step (list lookup, worktree setup) fails — the task is still Queued.
        var id = await SeedTaskAsync(TaskStatus.Queued);

        var result = await _sut.FailAsync(id, DateTime.UtcNow, "list not found", default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Failed, t.Status);
        Assert.Equal("list not found", t.Result);
    }

    [Fact]
    public async Task FailAsync_FromDone_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Done);

        var result = await _sut.FailAsync(id, DateTime.UtcNow, "boom", default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Done, await GetStatusAsync(id));
    }

    [Fact]
    public async Task FailAsync_FromWaitingForReview_IsNoOp()
    {
        var id = await SeedTaskAsync(TaskStatus.WaitingForReview);

        var result = await _sut.FailAsync(id, DateTime.UtcNow, "oops", default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.WaitingForReview, await GetStatusAsync(id));
    }

    // ─── CancelAsync ──────────────────────────────────────────────────────

    [Fact]
    public async Task CancelAsync_FromRunning_TransitionsToCancelled()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.CancelAsync(id, DateTime.UtcNow, default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Cancelled, await GetStatusAsync(id));
    }

    [Fact]
    public async Task CancelAsync_FromDone_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Done);

        var result = await _sut.CancelAsync(id, DateTime.UtcNow, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Done, await GetStatusAsync(id));
    }

    [Fact]
    public async Task CancelAsync_WhileUnitMergeDraining_Rejects_AndDoesNotMutate()
    {
        var mergeState = new FakeActiveMergeState();
        var built = TaskStateServiceBuilder.Build(_factory, () => mergeState);
        var id = await SeedTaskAsync(TaskStatus.WaitingForReview);
        mergeState.ActiveTaskIds.Add(id);

        var result = await built.State.CancelAsync(id, DateTime.UtcNow, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.WaitingForReview, await GetStatusAsync(id));
    }

    [Fact]
    public async Task CancelAsync_AfterUnitMergeDrainCompletes_Succeeds()
    {
        var mergeState = new FakeActiveMergeState();
        var built = TaskStateServiceBuilder.Build(_factory, () => mergeState);
        var id = await SeedTaskAsync(TaskStatus.WaitingForReview);
        // Merge drain finished (or never started) for this task — cancel behaves as before.

        var result = await built.State.CancelAsync(id, DateTime.UtcNow, default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Cancelled, await GetStatusAsync(id));
    }

    // ─── ResetToIdleAsync ─────────────────────────────────────────────────

    [Fact]
    public async Task ResetToIdleAsync_FromFailed_ClearsTimestamps()
    {
        var id = await SeedTaskAsync(TaskStatus.Failed);
        await using (var ctx = _factory.CreateDbContext())
        {
            await ctx.Tasks.Where(t => t.Id == id)
                .ExecuteUpdateAsync(s => s
                    .SetProperty(t => t.StartedAt, DateTime.UtcNow.AddMinutes(-5))
                    .SetProperty(t => t.FinishedAt, DateTime.UtcNow.AddMinutes(-1))
                    .SetProperty(t => t.Result, "old"));
        }

        var result = await _sut.ResetToIdleAsync(id, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Idle, t.Status);
        Assert.Null(t.StartedAt);
        Assert.Null(t.FinishedAt);
        Assert.Null(t.Result);
    }

    [Fact]
    public async Task ResetToIdleAsync_FromRunning_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.ResetToIdleAsync(id, default);

        Assert.False(result.Ok);
        Assert.Equal(TaskStatus.Running, await GetStatusAsync(id));
    }

    // ─── StartPlanningAsync ───────────────────────────────────────────────

    [Fact]
    public async Task StartPlanningAsync_FromIdle_SetsPlanningPhase()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle);

        var result = await _sut.StartPlanningAsync(id, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(TaskStatus.Idle, t.Status);
        Assert.Equal(PlanningPhase.Active, t.PlanningPhase);
    }

    [Fact]
    public async Task StartPlanningAsync_FromRunning_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Running);

        var result = await _sut.StartPlanningAsync(id, default);

        Assert.False(result.Ok);
    }

    // ─── FinalizePlanningAsync ────────────────────────────────────────────

    [Fact]
    public async Task FinalizePlanningAsync_OnActivePhase_TransitionsToFinalized()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle);
        await _sut.StartPlanningAsync(id, default);

        var result = await _sut.FinalizePlanningAsync(id, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(id);
        Assert.Equal(PlanningPhase.Finalized, t.PlanningPhase);
        Assert.NotNull(t.PlanningFinalizedAt);
    }

    [Fact]
    public async Task FinalizePlanningAsync_OnNonePhase_Rejects()
    {
        var id = await SeedTaskAsync(TaskStatus.Idle);

        var result = await _sut.FinalizePlanningAsync(id, default);

        Assert.False(result.Ok);
    }

    // ─── BlockOnAsync / UnblockAsync ─────────────────────────────────────

    [Fact]
    public async Task BlockOnAsync_SetsBlockedByTaskId()
    {
        var pred = await SeedTaskAsync(TaskStatus.Queued);
        var task = await SeedTaskAsync(TaskStatus.Queued);

        var result = await _sut.BlockOnAsync(task, pred, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(task);
        Assert.Equal(pred, t.BlockedByTaskId);
    }

    [Fact]
    public async Task UnblockAsync_ClearsBlockedByTaskId_AndWakesQueue()
    {
        var pred = await SeedTaskAsync(TaskStatus.Queued);
        var task = await SeedTaskAsync(TaskStatus.Queued);
        await _sut.BlockOnAsync(task, pred, default);
        var wakesBefore = _built.WakeCount();

        var result = await _sut.UnblockAsync(task, default);

        Assert.True(result.Ok);
        var t = await GetTaskAsync(task);
        Assert.Null(t.BlockedByTaskId);
        Assert.True(_built.WakeCount() > wakesBefore);
    }

    // ─── RecoverStaleRunningAsync ─────────────────────────────────────────

    [Fact]
    public async Task RecoverStaleRunningAsync_FlipsAllRunningToFailed_ReturnsCount()
    {
        var r1 = await SeedTaskAsync(TaskStatus.Running);
        var r2 = await SeedTaskAsync(TaskStatus.Running);
        var q = await SeedTaskAsync(TaskStatus.Queued);

        var count = await _sut.RecoverStaleRunningAsync("worker restart", default);

        Assert.Equal(2, count);
        Assert.Equal(TaskStatus.Failed, await GetStatusAsync(r1));
        Assert.Equal(TaskStatus.Failed, await GetStatusAsync(r2));
        Assert.Equal(TaskStatus.Queued, await GetStatusAsync(q));
        var t = await GetTaskAsync(r1);
        Assert.StartsWith("[stale] ", t.Result);
    }

    // ─── Child terminal → chain advance ───────────────────────────────────

    [Fact]
    public async Task CompleteAsync_OnChild_AdvancesNextBlockedSibling()
    {
        var parent = await SeedTaskAsync(TaskStatus.Idle, phase: PlanningPhase.Finalized);
        var c0 = await SeedTaskAsync(TaskStatus.Running, parentId: parent, sortOrder: 0);
        var c1 = await SeedTaskAsync(TaskStatus.Queued, parentId: parent, sortOrder: 1, blockedBy: c0);
        var c2 = await SeedTaskAsync(TaskStatus.Queued, parentId: parent, sortOrder: 2, blockedBy: c1);

        var result = await _sut.CompleteAsync(c0, DateTime.UtcNow, "ok", default);

        Assert.True(result.Ok);
        Assert.Equal(TaskStatus.Done, await GetStatusAsync(c0));
        // c1 was BlockedBy=c0 → chain coordinator unblocks → BlockedByTaskId cleared, still Queued.
        var t1 = await GetTaskAsync(c1);
        Assert.Equal(TaskStatus.Queued, t1.Status);
        Assert.Null(t1.BlockedByTaskId);
        // c2 still blocked on c1.
        var t2 = await GetTaskAsync(c2);
        Assert.Equal(TaskStatus.Queued, t2.Status);
        Assert.Equal(c1, t2.BlockedByTaskId);
    }
}