git.stella-ops.org/src/TaskRunner/StellaOps.TaskRunner/StellaOps.TaskRunner.Tests/MongoPackRunStoresTests.cs

using System.Text.Json.Nodes;
using Microsoft.Extensions.Logging.Abstractions;
using MongoDB.Driver;
using StellaOps.TaskRunner.Core.Execution;
using StellaOps.TaskRunner.Core.Execution.Simulation;
using StellaOps.TaskRunner.Core.Planning;
using StellaOps.TaskRunner.Core.TaskPacks;
using StellaOps.TaskRunner.Infrastructure.Execution;
using Xunit;
using Xunit.Sdk;

namespace StellaOps.TaskRunner.Tests;

public sealed class MongoPackRunStoresTests
{
    [Fact]
    public async Task StateStore_RoundTrips_State()
    {
        using var context = MongoTaskRunnerTestContext.Create();

        var mongoOptions = context.CreateMongoOptions();
        var stateStore = new MongoPackRunStateStore(context.Database, mongoOptions);

        var plan = CreatePlan();
        var executionContext = new PackRunExecutionContext("mongo-run-state", plan, DateTimeOffset.UtcNow);
        var graph = new PackRunExecutionGraphBuilder().Build(plan);
        var simulationEngine = new PackRunSimulationEngine();
        var state = PackRunStateFactory.CreateInitialState(executionContext, graph, simulationEngine, DateTimeOffset.UtcNow);

        await stateStore.SaveAsync(state, CancellationToken.None);

        var reloaded = await stateStore.GetAsync(state.RunId, CancellationToken.None);

        Assert.NotNull(reloaded);
        Assert.Equal(state.RunId, reloaded!.RunId);
        Assert.Equal(state.PlanHash, reloaded.PlanHash);
        Assert.Equal(state.Steps.Count, reloaded.Steps.Count);
    }

    [Fact]
    public async Task LogStore_Appends_And_Reads_In_Order()
    {
        using var context = MongoTaskRunnerTestContext.Create();
        var mongoOptions = context.CreateMongoOptions();
        var logStore = new MongoPackRunLogStore(context.Database, mongoOptions);

        var runId = "mongo-log";

        await logStore.AppendAsync(runId, new PackRunLogEntry(DateTimeOffset.UtcNow, "info", "run.created", "created", null, null), CancellationToken.None);
        await logStore.AppendAsync(runId, new PackRunLogEntry(DateTimeOffset.UtcNow.AddSeconds(1), "warn", "step.retry", "retry", "step-a", new Dictionary<string, string> { ["attempt"] = "2" }), CancellationToken.None);

        var entries = new List<PackRunLogEntry>();
        await foreach (var entry in logStore.ReadAsync(runId, CancellationToken.None))
        {
            entries.Add(entry);
        }

        Assert.Equal(2, entries.Count);
        Assert.Equal("run.created", entries[0].EventType);
        Assert.Equal("step.retry", entries[1].EventType);
        Assert.Equal("step-a", entries[1].StepId);
        Assert.True(await logStore.ExistsAsync(runId, CancellationToken.None));
    }

    [Fact]
    public async Task ApprovalStore_RoundTrips_And_Updates()
    {
        using var context = MongoTaskRunnerTestContext.Create();
        var mongoOptions = context.CreateMongoOptions();
        var approvalStore = new MongoPackRunApprovalStore(context.Database, mongoOptions);

        var runId = "mongo-approvals";
        var approval = new PackRunApprovalState(
            "security-review",
            new[] { "packs.approve" },
            new[] { "step-plan" },
            Array.Empty<string>(),
            reasonTemplate: "Security approval required.",
            DateTimeOffset.UtcNow,
            PackRunApprovalStatus.Pending);

        await approvalStore.SaveAsync(runId, new[] { approval }, CancellationToken.None);

        var approvals = await approvalStore.GetAsync(runId, CancellationToken.None);
        Assert.Single(approvals);

        var updated = approval.Approve("approver", DateTimeOffset.UtcNow, "Approved");
        await approvalStore.UpdateAsync(runId, updated, CancellationToken.None);

        approvals = await approvalStore.GetAsync(runId, CancellationToken.None);
        Assert.Single(approvals);
        Assert.Equal(PackRunApprovalStatus.Approved, approvals[0].Status);
        Assert.Equal("approver", approvals[0].ActorId);
    }

    [Fact]
    public async Task ArtifactUploader_Persists_Metadata()
    {
        using var context = MongoTaskRunnerTestContext.Create();
        var mongoOptions = context.CreateMongoOptions();
        var database = context.Database;

        var artifactUploader = new MongoPackRunArtifactUploader(
            database,
            mongoOptions,
            TimeProvider.System,
            NullLogger<MongoPackRunArtifactUploader>.Instance);

        var plan = CreatePlanWithOutputs(out var outputFile);
        try
        {
            var executionContext = new PackRunExecutionContext("mongo-artifacts", plan, DateTimeOffset.UtcNow);
            var graph = new PackRunExecutionGraphBuilder().Build(plan);
            var simulationEngine = new PackRunSimulationEngine();
            var state = PackRunStateFactory.CreateInitialState(executionContext, graph, simulationEngine, DateTimeOffset.UtcNow);

            await artifactUploader.UploadAsync(executionContext, state, plan.Outputs, CancellationToken.None);

            var documents = await database
                .GetCollection<MongoPackRunArtifactUploader.PackRunArtifactDocument>(mongoOptions.ArtifactsCollection)
                .Find(Builders<MongoPackRunArtifactUploader.PackRunArtifactDocument>.Filter.Empty)
                .ToListAsync(TestContext.Current.CancellationToken);

            var bundleDocument = Assert.Single(documents, d => string.Equals(d.Name, "bundlePath", StringComparison.Ordinal));
            Assert.Equal("file", bundleDocument.Type);
            Assert.Equal(outputFile, bundleDocument.SourcePath);
            Assert.Equal("referenced", bundleDocument.Status);
        }
        finally
        {
            if (File.Exists(outputFile))
            {
                File.Delete(outputFile);
            }
        }
    }

    private static TaskPackPlan CreatePlan()
    {
        var manifest = TestManifests.Load(TestManifests.Sample);
        var planner = new TaskPackPlanner();
        var result = planner.Plan(manifest);
        if (!result.Success || result.Plan is null)
        {
            Assert.Skip("Failed to build task pack plan for Mongo tests.");
            throw new InvalidOperationException();
        }

        return result.Plan;
    }

    private static TaskPackPlan CreatePlanWithOutputs(out string outputFile)
    {
        var manifest = TestManifests.Load(TestManifests.Output);
        var planner = new TaskPackPlanner();
        var result = planner.Plan(manifest);
        if (!result.Success || result.Plan is null)
        {
            Assert.Skip("Failed to build output plan for Mongo tests.");
            throw new InvalidOperationException();
        }

        // Materialize a fake output file referenced by the plan.
        outputFile = Path.Combine(Path.GetTempPath(), $"taskrunner-output-{Guid.NewGuid():N}.txt");
        File.WriteAllText(outputFile, "fixture");

        // Update the plan output path parameter to point at the file we just created.
        var originalPlan = result.Plan;

        var resolvedFile = outputFile;

        var outputs = originalPlan.Outputs
            .Select(output =>
            {
                if (!string.Equals(output.Name, "bundlePath", StringComparison.Ordinal))
                {
                    return output;
                }

                var node = JsonNode.Parse($"\"{resolvedFile.Replace("\\", "\\\\")}\"");
                var parameter = new TaskPackPlanParameterValue(node, null, null, false);
                return output with { Path = parameter };
            })
            .ToArray();

        return new TaskPackPlan(
            originalPlan.Metadata,
            originalPlan.Inputs,
            originalPlan.Steps,
            originalPlan.Hash,
            originalPlan.Approvals,
            originalPlan.Secrets,
            outputs,
            originalPlan.FailurePolicy);
    }
}