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

// Unit tests for MQTT will message delivery on abnormal disconnection.
// Go reference: golang/nats-server/server/mqtt.go — mqttDeliverWill ~line 490,
//   TestMQTTWill server/mqtt_test.go:4129

using NATS.Server.Mqtt;
using Shouldly;

namespace NATS.Server.Mqtt.Tests.Mqtt;

public class MqttWillMessageTests
{
    // Go ref: mqtt.go mqttSession will field — will message is stored on CONNECT with will flag.
    [Fact]
    public void SetWill_stores_will_message()
    {
        var store = new MqttSessionStore();
        var will = new WillMessage { Topic = "client/status", Payload = "offline"u8.ToArray(), QoS = 0, Retain = false };

        store.SetWill("client-1", will);

        var stored = store.GetWill("client-1");
        stored.ShouldNotBeNull();
        stored.Topic.ShouldBe("client/status");
        stored.Payload.ShouldBe("offline"u8.ToArray());
    }

    // Go ref: mqttDeliverWill — on graceful DISCONNECT, will is cleared (not delivered).
    [Fact]
    public void ClearWill_removes_will()
    {
        var store = new MqttSessionStore();
        var will = new WillMessage { Topic = "client/status", Payload = "offline"u8.ToArray() };

        store.SetWill("client-2", will);
        store.ClearWill("client-2");

        store.GetWill("client-2").ShouldBeNull();
    }

    // Go ref: TestMQTTWill server/mqtt_test.go:4129 — will is published on abnormal disconnect.
    [Fact]
    public void PublishWillMessage_publishes_on_abnormal_disconnect()
    {
        var store = new MqttSessionStore();
        string? publishedTopic = null;
        byte[]? publishedPayload = null;
        byte publishedQoS = 0xFF;
        bool publishedRetain = false;

        store.OnPublish = (topic, payload, qos, retain) =>
        {
            publishedTopic = topic;
            publishedPayload = payload;
            publishedQoS = qos;
            publishedRetain = retain;
        };

        var will = new WillMessage { Topic = "device/gone", Payload = "disconnected"u8.ToArray(), QoS = 1, Retain = false };
        store.SetWill("client-3", will);

        var result = store.PublishWillMessage("client-3");

        result.ShouldBeTrue();
        publishedTopic.ShouldBe("device/gone");
        publishedPayload.ShouldBe("disconnected"u8.ToArray());
        publishedQoS.ShouldBe((byte)1);
        publishedRetain.ShouldBeFalse();
    }

    // Go ref: mqttDeliverWill — no-op when no will is registered.
    [Fact]
    public void PublishWillMessage_returns_false_when_no_will()
    {
        var store = new MqttSessionStore();
        var invoked = false;
        store.OnPublish = (_, _, _, _) => { invoked = true; };

        var result = store.PublishWillMessage("client-no-will");

        result.ShouldBeFalse();
        invoked.ShouldBeFalse();
    }

    // Go ref: mqttDeliverWill — will is consumed (not published twice).
    [Fact]
    public void PublishWillMessage_clears_will_after_publish()
    {
        var store = new MqttSessionStore();
        store.OnPublish = (_, _, _, _) => { };

        var will = new WillMessage { Topic = "sensor/status", Payload = "gone"u8.ToArray() };
        store.SetWill("client-5", will);

        store.PublishWillMessage("client-5");

        store.GetWill("client-5").ShouldBeNull();
        store.PublishWillMessage("client-5").ShouldBeFalse();
    }

    // Go ref: TestMQTTWill — graceful DISCONNECT clears the will before disconnect;
    //   subsequent PublishWillMessage has no effect.
    [Fact]
    public void CleanDisconnect_does_not_publish_will()
    {
        var store = new MqttSessionStore();
        var invoked = false;
        store.OnPublish = (_, _, _, _) => { invoked = true; };

        var will = new WillMessage { Topic = "client/status", Payload = "bye"u8.ToArray() };
        store.SetWill("client-6", will);

        // Simulate graceful DISCONNECT: clear will before triggering publish path
        store.ClearWill("client-6");
        var result = store.PublishWillMessage("client-6");

        result.ShouldBeFalse();
        invoked.ShouldBeFalse();
    }

    // Go ref: TestMQTTWill — published topic and payload must exactly match what was registered.
    [Fact]
    public void WillMessage_preserves_topic_and_payload()
    {
        var store = new MqttSessionStore();
        var capturedTopic = string.Empty;
        var capturedPayload = Array.Empty<byte>();
        store.OnPublish = (topic, payload, _, _) =>
        {
            capturedTopic = topic;
            capturedPayload = payload;
        };

        var originalPayload = "sensor-offline-payload"u8.ToArray();
        store.SetWill("client-7", new WillMessage { Topic = "sensors/temperature/offline", Payload = originalPayload });
        store.PublishWillMessage("client-7");

        capturedTopic.ShouldBe("sensors/temperature/offline");
        capturedPayload.ShouldBe(originalPayload);
    }

    // Go ref: TestMQTTWill — QoS level from the will is forwarded to the broker publish path.
    [Fact]
    public void WillMessage_preserves_qos()
    {
        var store = new MqttSessionStore();
        byte capturedQoS = 0xFF;
        store.OnPublish = (_, _, qos, _) => { capturedQoS = qos; };

        store.SetWill("client-8", new WillMessage { Topic = "t", Payload = [], QoS = 1 });
        store.PublishWillMessage("client-8");

        capturedQoS.ShouldBe((byte)1);
    }

    // Go ref: TestMQTTWillRetain — retain flag from the will is forwarded to the broker publish path.
    [Fact]
    public void WillMessage_preserves_retain_flag()
    {
        var store = new MqttSessionStore();
        bool capturedRetain = false;
        store.OnPublish = (_, _, _, retain) => { capturedRetain = retain; };

        store.SetWill("client-9", new WillMessage { Topic = "t", Payload = [], Retain = true });
        store.PublishWillMessage("client-9");

        capturedRetain.ShouldBeTrue();
    }

    // Go ref: MQTT 5.0 Will-Delay-Interval — a will with delay > 0 is not immediately published;
    //   it is tracked as a delayed will and OnPublish is NOT called immediately.
    [Fact]
    public void PublishWillMessage_with_delay_stores_delayed_will_and_does_not_call_OnPublish()
    {
        var store = new MqttSessionStore();
        var immediatelyPublished = false;
        store.OnPublish = (_, _, _, _) => { immediatelyPublished = true; };

        var will = new WillMessage
        {
            Topic = "device/status",
            Payload = "gone"u8.ToArray(),
            QoS = 0,
            Retain = false,
            DelayIntervalSeconds = 30
        };
        store.SetWill("client-10", will);

        var result = store.PublishWillMessage("client-10");

        // Returns true because a will was found
        result.ShouldBeTrue();

        // OnPublish must NOT have been called — it is delayed
        immediatelyPublished.ShouldBeFalse();

        // The will must be tracked as a pending delayed will
        var delayed = store.GetDelayedWill("client-10");
        delayed.ShouldNotBeNull();
        delayed!.Value.Will.Topic.ShouldBe("device/status");
        delayed.Value.Will.DelayIntervalSeconds.ShouldBe(30);
    }
}