git.stella-ops.org/tests/AirGap/StellaOps.AirGap.Controller.Tests/InMemoryAirGapStateStoreTests.cs

using StellaOps.AirGap.Controller.Domain;
using StellaOps.AirGap.Controller.Stores;
using StellaOps.AirGap.Time.Models;
using Xunit;

namespace StellaOps.AirGap.Controller.Tests;

public class InMemoryAirGapStateStoreTests
{
    private readonly InMemoryAirGapStateStore _store = new();

    [Fact]
    public async Task Upsert_and_read_state_by_tenant()
    {
        var state = new AirGapState
        {
            TenantId = "tenant-x",
            Sealed = true,
            PolicyHash = "hash-1",
            TimeAnchor = new TimeAnchor(DateTimeOffset.UtcNow, "roughtime", "roughtime", "fp", "digest"),
            StalenessBudget = new StalenessBudget(10, 20),
            LastTransitionAt = DateTimeOffset.UtcNow
        };

        await _store.SetAsync(state);

        var stored = await _store.GetAsync("tenant-x");
        Assert.True(stored.Sealed);
        Assert.Equal("hash-1", stored.PolicyHash);
        Assert.Equal("tenant-x", stored.TenantId);
        Assert.Equal(state.TimeAnchor.TokenDigest, stored.TimeAnchor.TokenDigest);
        Assert.Equal(10, stored.StalenessBudget.WarningSeconds);
    }

    [Fact]
    public async Task Enforces_singleton_per_tenant()
    {
        var first = new AirGapState { TenantId = "tenant-y", Sealed = true, PolicyHash = "h1" };
        var second = new AirGapState { TenantId = "tenant-y", Sealed = false, PolicyHash = "h2" };

        await _store.SetAsync(first);
        await _store.SetAsync(second);

        var stored = await _store.GetAsync("tenant-y");
        Assert.Equal("h2", stored.PolicyHash);
        Assert.False(stored.Sealed);
    }

    [Fact]
    public async Task Defaults_to_unknown_when_missing()
    {
        var stored = await _store.GetAsync("absent");
        Assert.False(stored.Sealed);
        Assert.Equal("absent", stored.TenantId);
    }

    [Fact]
    public async Task Parallel_upserts_keep_single_document()
    {
        var tasks = Enumerable.Range(0, 20).Select(i =>
        {
            var state = new AirGapState
            {
                TenantId = "tenant-parallel",
                Sealed = i % 2 == 0,
                PolicyHash = $"hash-{i}"
            };
            return _store.SetAsync(state);
        });

        await Task.WhenAll(tasks);

        var stored = await _store.GetAsync("tenant-parallel");
        Assert.StartsWith("hash-", stored.PolicyHash);
    }

    [Fact]
    public async Task Multi_tenant_updates_do_not_collide()
    {
        var tenants = Enumerable.Range(0, 5).Select(i => $"t-{i}").ToArray();

        var tasks = tenants.Select(t => _store.SetAsync(new AirGapState
        {
            TenantId = t,
            Sealed = true,
            PolicyHash = $"hash-{t}"
        }));

        await Task.WhenAll(tasks);

        foreach (var t in tenants)
        {
            var stored = await _store.GetAsync(t);
            Assert.Equal($"hash-{t}", stored.PolicyHash);
        }
    }

    [Fact]
    public async Task Staleness_round_trip_matches_budget()
    {
        var anchor = new TimeAnchor(DateTimeOffset.UtcNow.AddMinutes(-3), "roughtime", "roughtime", "fp", "digest");
        var budget = new StalenessBudget(60, 600);
        await _store.SetAsync(new AirGapState
        {
            TenantId = "tenant-staleness",
            Sealed = true,
            PolicyHash = "hash-s",
            TimeAnchor = anchor,
            StalenessBudget = budget,
            LastTransitionAt = DateTimeOffset.UtcNow
        });

        var stored = await _store.GetAsync("tenant-staleness");
        Assert.Equal(anchor.TokenDigest, stored.TimeAnchor.TokenDigest);
        Assert.Equal(budget.WarningSeconds, stored.StalenessBudget.WarningSeconds);
        Assert.Equal(budget.BreachSeconds, stored.StalenessBudget.BreachSeconds);
    }

    [Fact]
    public async Task Multi_tenant_states_preserve_transition_times()
    {
        var tenants = new[] { "a", "b", "c" };
        var now = DateTimeOffset.UtcNow;

        foreach (var t in tenants)
        {
            await _store.SetAsync(new AirGapState
            {
                TenantId = t,
                Sealed = true,
                PolicyHash = $"ph-{t}",
                LastTransitionAt = now
            });
        }

        foreach (var t in tenants)
        {
            var state = await _store.GetAsync(t);
            Assert.Equal(now, state.LastTransitionAt);
            Assert.Equal($"ph-{t}", state.PolicyHash);
        }
    }
}