diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/StoreSessionReregistrationPreserveTests.cs b/src/Backend/AutopilotMonitor.Functions.Tests/StoreSessionReregistrationPreserveTests.cs new file mode 100644 index 00000000..fee7f198 --- /dev/null +++ b/src/Backend/AutopilotMonitor.Functions.Tests/StoreSessionReregistrationPreserveTests.cs @@ -0,0 +1,151 @@ +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Azure; +using Azure.Data.Tables; +using AutopilotMonitor.Functions.Services; +using AutopilotMonitor.Shared; +using AutopilotMonitor.Shared.Models; +using Microsoft.Extensions.Logging.Abstractions; +using Moq; +using Xunit; + +namespace AutopilotMonitor.Functions.Tests; + +/// +/// Re-registration preservation contract for . +/// On agent restart the agent re-registers the SAME session id; StoreSessionAsync rewrites the +/// Sessions row with UpsertEntity (Replace mode), which drops any field NOT carried over from the +/// existing row. Cumulative counters maintained incrementally by IncrementSessionEventCountAsync +/// (EventCount, PlatformScriptCount, RemediationScriptCount) must therefore be preserved — otherwise +/// every restart zeroes them, undercounting scripts and drifting Sessions vs. SessionsIndex (the +/// index keeps the pre-restart high-water value because the increment-merge only writes a counter +/// when its increment > 0). Regression guard for that drift. +/// +public class StoreSessionReregistrationPreserveTests +{ + private const string TenantId = "11111111-1111-1111-1111-111111111111"; + private const string SessionId = "22222222-2222-2222-2222-222222222222"; + + [Fact] + public async Task StoreSessionAsync_preserves_cumulative_counters_across_reregistration() + { + // Existing row as left by a prior agent run: counters already accumulated. + var existing = new TableEntity(TenantId, SessionId) + { + ["StartedAt"] = new DateTimeOffset(new DateTime(2026, 1, 1, 8, 0, 0, DateTimeKind.Utc)), + ["Status"] = "InProgress", + ["EventCount"] = 42, + ["PlatformScriptCount"] = 3, + ["RemediationScriptCount"] = 5, + }; + existing.ETag = new ETag("0xEXISTING"); + + var harness = new Harness(existing); + + var registration = new SessionRegistration + { + TenantId = TenantId, + SessionId = SessionId, + // Agent restart reports a later StartedAt — StoreSessionAsync keeps the earlier existing one, + // but that is orthogonal to the counter-preservation under test. + StartedAt = new DateTime(2026, 1, 1, 9, 0, 0, DateTimeKind.Utc), + SerialNumber = "SN-1", + Manufacturer = "Contoso", + Model = "Model-X", + DeviceName = "PC-1", + }; + + var ok = await harness.Sut.StoreSessionAsync(registration); + + Assert.True(ok); + Assert.NotNull(harness.Written); + // The Replace-written row must still carry the cumulative counters from the existing row. + Assert.Equal(3, harness.Written!.GetInt32("PlatformScriptCount")); + Assert.Equal(5, harness.Written.GetInt32("RemediationScriptCount")); + // EventCount preservation is the established sibling behaviour — assert it as a sanity anchor. + Assert.Equal(42, harness.Written.GetInt32("EventCount")); + } + + [Fact] + public async Task StoreSessionAsync_defaults_counters_to_zero_for_a_brand_new_session() + { + // No existing row (first registration) → counters default to 0, never absent. + var harness = new Harness(existing: null); + + var registration = new SessionRegistration + { + TenantId = TenantId, + SessionId = SessionId, + StartedAt = new DateTime(2026, 1, 1, 9, 0, 0, DateTimeKind.Utc), + SerialNumber = "SN-1", + Manufacturer = "Contoso", + Model = "Model-X", + DeviceName = "PC-1", + }; + + var ok = await harness.Sut.StoreSessionAsync(registration); + + Assert.True(ok); + Assert.NotNull(harness.Written); + Assert.Equal(0, harness.Written!.GetInt32("PlatformScriptCount")); + Assert.Equal(0, harness.Written.GetInt32("RemediationScriptCount")); + } + + // ============================================================ Harness ==== + + /// + /// SDK-mock harness. Only the Sessions table client is configured; the Events query + /// (earliest-timestamp probe) and the SessionsIndex upsert are both wrapped in swallow-all + /// try/catch in the SUT, so leaving those table clients unconfigured (null) is harmless and + /// keeps the test focused on the Sessions Replace payload. + /// + private sealed class Harness + { + public Mock Sessions { get; } + public TableStorageService Sut { get; } + public TableEntity? Written { get; private set; } + + public Harness(TableEntity? existing) + { + Sessions = new Mock(); + + if (existing is not null) + { + Sessions.Setup(t => t.GetEntityAsync( + TenantId, SessionId, It.IsAny>(), It.IsAny())) + .ReturnsAsync(Response.FromValue(existing, new Mock().Object)); + } + else + { + Sessions.Setup(t => t.GetEntityAsync( + TenantId, SessionId, It.IsAny>(), It.IsAny())) + .ThrowsAsync(new RequestFailedException(404, "ResourceNotFound")); + } + + // Replace write (existing row) — capture the payload. + Sessions.Setup(t => t.UpdateEntityAsync( + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny())) + .Returns((e, _, _, _) => + { + Written = e; + return Task.FromResult(new Mock().Object); + }); + + // Add write (brand-new row, no ETag) — capture the payload. + Sessions.Setup(t => t.AddEntityAsync(It.IsAny(), It.IsAny())) + .Returns((e, _) => + { + Written = e; + return Task.FromResult(new Mock().Object); + }); + + var mockServiceClient = new Mock(); + mockServiceClient.Setup(s => s.GetTableClient(Constants.TableNames.Sessions)).Returns(Sessions.Object); + // Events / SessionsIndex intentionally unconfigured → null client → SUT swallows. + + Sut = new TableStorageService(mockServiceClient.Object, NullLogger.Instance); + } + } +} diff --git a/src/Backend/AutopilotMonitor.Functions/Services/TableStorageService.Sessions.cs b/src/Backend/AutopilotMonitor.Functions/Services/TableStorageService.Sessions.cs index d262e353..c4e4e0db 100644 --- a/src/Backend/AutopilotMonitor.Functions/Services/TableStorageService.Sessions.cs +++ b/src/Backend/AutopilotMonitor.Functions/Services/TableStorageService.Sessions.cs @@ -334,6 +334,11 @@ namespace AutopilotMonitor.Functions.Services int currentPhase = (int)EnrollmentPhase.Start; string status = SessionStatus.InProgress.ToString(); int eventCount = 0; + // Cumulative script counters are maintained incrementally by IncrementSessionEventCountAsync. + // They MUST survive re-registration — the Replace below would otherwise zero them on every + // agent restart, undercounting scripts run before the restart and drifting Sessions vs. SessionsIndex. + int platformScriptCount = 0; + int remediationScriptCount = 0; DateTime? completedAt = null; string failureReason = string.Empty; bool isPreProvisioned = registration.IsPreProvisioned; @@ -367,6 +372,8 @@ namespace AutopilotMonitor.Functions.Services currentPhase = existingEntity.GetInt32("CurrentPhase") ?? currentPhase; status = existingEntity.GetString("Status") ?? status; eventCount = existingEntity.GetInt32("EventCount") ?? eventCount; + platformScriptCount = existingEntity.GetInt32("PlatformScriptCount") ?? platformScriptCount; + remediationScriptCount = existingEntity.GetInt32("RemediationScriptCount") ?? remediationScriptCount; completedAt = existingEntity.GetDateTimeOffset("CompletedAt")?.UtcDateTime; failureReason = existingEntity.GetString("FailureReason") ?? string.Empty; @@ -442,7 +449,9 @@ namespace AutopilotMonitor.Functions.Services ["EnrollmentType"] = registration.EnrollmentType ?? "v1", ["CurrentPhase"] = currentPhase, ["Status"] = status, - ["EventCount"] = eventCount + ["EventCount"] = eventCount, + ["PlatformScriptCount"] = platformScriptCount, + ["RemediationScriptCount"] = remediationScriptCount }; if (completedAt.HasValue)