up

tests/AirGap/StellaOps.AirGap.Controller.Tests/MongoAirGapStateStoreTests.cs

@@ -1,3 +1,4 @@
+using MongoDB.Bson;
 using MongoDB.Driver;
 using StellaOps.AirGap.Controller.Domain;
 using StellaOps.AirGap.Controller.Stores;
@@ -66,6 +67,115 @@ public class MongoAirGapStateStoreTests : IDisposable
         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();