natsdotnet/tests/NATS.Server.Mqtt.Tests/Mqtt/MqttSessionPersistenceTests.cs

// MQTT session persistence tests.
// Go reference: golang/nats-server/server/mqtt.go:253-360
//   Session store          — mqttInitSessionStore / mqttStoreSession / mqttLoadSession
//   Flapper detection      — mqttCheckFlapper (~lines 300–360)

using NATS.Server.Mqtt;

namespace NATS.Server.Mqtt.Tests.Mqtt;

public class MqttSessionPersistenceTests
{
    [Fact]
    public void SaveSession_ThenLoad_RoundTrips()
    {
        // Go reference: server/mqtt.go mqttStoreSession / mqttLoadSession
        var store = new MqttSessionStore();
        var session = new MqttSessionData
        {
            ClientId = "client-1",
            Subscriptions = new Dictionary<string, int> { ["sensors/temp"] = 1, ["alerts/#"] = 0 },
            PendingPublishes = [new MqttPendingPublish(42, "sensors/temp", "72.5")],
            WillTopic = "clients/offline",
            WillPayload = [0x01, 0x02],
            WillQoS = 1,
            WillRetain = true,
            CleanSession = false,
            ConnectedAtUtc = new DateTime(2026, 1, 15, 10, 30, 0, DateTimeKind.Utc),
            LastActivityUtc = new DateTime(2026, 1, 15, 10, 35, 0, DateTimeKind.Utc),
        };

        store.SaveSession(session);
        var loaded = store.LoadSession("client-1");

        loaded.ShouldNotBeNull();
        loaded.ClientId.ShouldBe("client-1");
        loaded.Subscriptions.Count.ShouldBe(2);
        loaded.Subscriptions["sensors/temp"].ShouldBe(1);
        loaded.Subscriptions["alerts/#"].ShouldBe(0);
        loaded.PendingPublishes.Count.ShouldBe(1);
        loaded.PendingPublishes[0].PacketId.ShouldBe(42);
        loaded.PendingPublishes[0].Topic.ShouldBe("sensors/temp");
        loaded.PendingPublishes[0].Payload.ShouldBe("72.5");
        loaded.WillTopic.ShouldBe("clients/offline");
        loaded.WillPayload.ShouldBe(new byte[] { 0x01, 0x02 });
        loaded.WillQoS.ShouldBe(1);
        loaded.WillRetain.ShouldBeTrue();
        loaded.CleanSession.ShouldBeFalse();
        loaded.ConnectedAtUtc.ShouldBe(new DateTime(2026, 1, 15, 10, 30, 0, DateTimeKind.Utc));
        loaded.LastActivityUtc.ShouldBe(new DateTime(2026, 1, 15, 10, 35, 0, DateTimeKind.Utc));
    }

    [Fact]
    public void SaveSession_Update_OverwritesPrevious()
    {
        // Go reference: server/mqtt.go mqttStoreSession — overwrites existing
        var store = new MqttSessionStore();

        store.SaveSession(new MqttSessionData
        {
            ClientId = "client-x",
            Subscriptions = new Dictionary<string, int> { ["old/topic"] = 0 },
        });

        store.SaveSession(new MqttSessionData
        {
            ClientId = "client-x",
            Subscriptions = new Dictionary<string, int> { ["new/topic"] = 1 },
        });

        var loaded = store.LoadSession("client-x");
        loaded.ShouldNotBeNull();
        loaded.Subscriptions.ShouldContainKey("new/topic");
        loaded.Subscriptions.ShouldNotContainKey("old/topic");
    }

    [Fact]
    public void LoadSession_NonExistent_ReturnsNull()
    {
        // Go reference: server/mqtt.go mqttLoadSession — returns nil for missing
        var store = new MqttSessionStore();

        var loaded = store.LoadSession("does-not-exist");

        loaded.ShouldBeNull();
    }

    [Fact]
    public void DeleteSession_RemovesFromStore()
    {
        // Go reference: server/mqtt.go mqttDeleteSession
        var store = new MqttSessionStore();
        store.SaveSession(new MqttSessionData { ClientId = "to-delete" });

        store.DeleteSession("to-delete");

        store.LoadSession("to-delete").ShouldBeNull();
    }

    [Fact]
    public void DeleteSession_NonExistent_NoError()
    {
        // Go reference: server/mqtt.go mqttDeleteSession — no-op on missing
        var store = new MqttSessionStore();

        // Should not throw
        store.DeleteSession("phantom");

        store.LoadSession("phantom").ShouldBeNull();
    }

    [Fact]
    public void ListSessions_ReturnsAllActive()
    {
        // Go reference: server/mqtt.go session enumeration
        var store = new MqttSessionStore();
        store.SaveSession(new MqttSessionData { ClientId = "alpha" });
        store.SaveSession(new MqttSessionData { ClientId = "beta" });
        store.SaveSession(new MqttSessionData { ClientId = "gamma" });

        var sessions = store.ListSessions();

        sessions.Count.ShouldBe(3);
        sessions.Select(s => s.ClientId).ShouldBe(
            new[] { "alpha", "beta", "gamma" },
            ignoreOrder: true);
    }

    [Fact]
    public void FlapperDetection_ThreeConnectsInTenSeconds_BackoffApplied()
    {
        // Go reference: server/mqtt.go mqttCheckFlapper ~line 300
        // Three connects within the flap window triggers backoff.
        var fakeTime = new FakeTimeProvider(new DateTimeOffset(2026, 1, 15, 12, 0, 0, TimeSpan.Zero));
        var store = new MqttSessionStore(
            flapWindow: TimeSpan.FromSeconds(10),
            flapThreshold: 3,
            flapBackoff: TimeSpan.FromSeconds(1),
            timeProvider: fakeTime);

        // Three rapid connects
        store.TrackConnectDisconnect("flapper", connected: true);
        fakeTime.Advance(TimeSpan.FromSeconds(1));
        store.TrackConnectDisconnect("flapper", connected: true);
        fakeTime.Advance(TimeSpan.FromSeconds(1));
        store.TrackConnectDisconnect("flapper", connected: true);

        var backoff = store.ShouldApplyBackoff("flapper");
        backoff.ShouldBeGreaterThan(TimeSpan.Zero);
        backoff.ShouldBe(TimeSpan.FromSeconds(1));
    }

    [Fact]
    public void FlapperDetection_SlowConnects_NoBackoff()
    {
        // Go reference: server/mqtt.go mqttCheckFlapper — slow connects should not trigger
        var fakeTime = new FakeTimeProvider(new DateTimeOffset(2026, 1, 15, 12, 0, 0, TimeSpan.Zero));
        var store = new MqttSessionStore(
            flapWindow: TimeSpan.FromSeconds(10),
            flapThreshold: 3,
            flapBackoff: TimeSpan.FromSeconds(1),
            timeProvider: fakeTime);

        // Three connects, but spread out beyond the window
        store.TrackConnectDisconnect("slow-client", connected: true);
        fakeTime.Advance(TimeSpan.FromSeconds(5));
        store.TrackConnectDisconnect("slow-client", connected: true);
        fakeTime.Advance(TimeSpan.FromSeconds(6)); // first connect now outside window
        store.TrackConnectDisconnect("slow-client", connected: true);

        var backoff = store.ShouldApplyBackoff("slow-client");
        backoff.ShouldBe(TimeSpan.Zero);
    }

    [Fact]
    public void CleanSession_DeletesOnConnect()
    {
        // Go reference: server/mqtt.go — clean session flag clears stored state
        var store = new MqttSessionStore();

        // Pre-populate a session
        store.SaveSession(new MqttSessionData
        {
            ClientId = "ephemeral",
            Subscriptions = new Dictionary<string, int> { ["topic/a"] = 1 },
            CleanSession = false,
        });

        store.LoadSession("ephemeral").ShouldNotBeNull();

        // Simulate clean session connect: delete the old session
        store.DeleteSession("ephemeral");

        store.LoadSession("ephemeral").ShouldBeNull();
    }
}

/// <summary>
/// Fake <see cref="TimeProvider"/> for deterministic time control in tests.
/// </summary>
internal sealed class FakeTimeProvider(DateTimeOffset startTime) : TimeProvider
{
    private DateTimeOffset _current = startTime;

    public override DateTimeOffset GetUtcNow() => _current;

    public void Advance(TimeSpan duration) => _current += duration;

    public void SetUtcNow(DateTimeOffset value) => _current = value;
}