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

using MongoDB.Bson;
using MongoDB.Driver;
using StellaOps.AirGap.Controller.Domain;
using StellaOps.AirGap.Controller.Stores;
using StellaOps.AirGap.Time.Models;
using StellaOps.Testing;
using Xunit;

namespace StellaOps.AirGap.Controller.Tests;

public class MongoAirGapStateStoreTests : IDisposable
{
    private readonly MongoRunnerFixture _mongo = new();
    private readonly IMongoCollection<AirGapStateDocument> _collection;
    private readonly MongoAirGapStateStore _store;

    public MongoAirGapStateStoreTests()
    {
        OpenSslAutoInit.Init();
        var database = _mongo.Client.GetDatabase("airgap_tests");
        _collection = MongoAirGapStateStore.EnsureCollection(database);
        _store = new MongoAirGapStateStore(_collection);
    }

    [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 Creates_unique_index_on_tenant_and_id()
    {
        var indexes = await _collection.Indexes.List().ToListAsync();
        var match = indexes.FirstOrDefault(idx =>
        {
            var key = idx["key"].AsBsonDocument;
            return key.ElementCount == 2
                   && key.Names.ElementAt(0) == "tenant_id"
                   && key.Names.ElementAt(1) == "_id";
        });

        Assert.NotNull(match);
        Assert.True(match!["unique"].AsBoolean);
    }

    [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);

        var count = await _collection.CountDocumentsAsync(Builders<AirGapStateDocument>.Filter.Eq(x => x.TenantId, "tenant-parallel"));
        Assert.Equal(1, count);
    }

    [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);
        }

        var totalDocs = await _collection.CountDocumentsAsync(FilterDefinition<AirGapStateDocument>.Empty);
        Assert.Equal(tenants.Length, totalDocs);
    }

    [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);
        }
    }

    public void Dispose()
    {
        _mongo.Dispose();
    }
}