CBDDC/tests/ZB.MOM.WW.CBDDC.Network.Tests/ProtocolTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using Google.Protobuf;
using ZB.MOM.WW.CBDDC.Core;
using ZB.MOM.WW.CBDDC.Network.Proto;
using ZB.MOM.WW.CBDDC.Network.Protocol;
using ZB.MOM.WW.CBDDC.Network.Security;

namespace ZB.MOM.WW.CBDDC.Network.Tests;

public class ProtocolTests
{
    private readonly ProtocolHandler _handler;

    /// <summary>
    ///     Initializes a new instance of the <see cref="ProtocolTests" /> class.
    /// </summary>
    public ProtocolTests()
    {
        _handler = new ProtocolHandler(NullLogger<ProtocolHandler>.Instance);
    }

    /// <summary>
    ///     Verifies a plain message can be written and read without transformation.
    /// </summary>
    [Fact]
    public async Task RoundTrip_ShouldWorks_WithPlainMessage()
    {
        // Arrange
        var stream = new MemoryStream();
        var message = new HandshakeRequest { NodeId = "node-1", AuthToken = "token" };

        // Act
        await _handler.SendMessageAsync(stream, MessageType.HandshakeReq, message, false, null);

        stream.Position = 0; // Reset for reading
        (var type, byte[] payload) = await _handler.ReadMessageAsync(stream, null);

        // Assert
        type.ShouldBe(MessageType.HandshakeReq);
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);
        decoded.NodeId.ShouldBe("node-1");
        decoded.AuthToken.ShouldBe("token");
    }

    /// <summary>
    ///     Verifies a compressed message can be written and read successfully.
    /// </summary>
    [Fact]
    public async Task RoundTrip_ShouldWork_WithCompression()
    {
        // Arrange
        var stream = new MemoryStream();
        // Create a large message to trigger compression logic (threshold is small but let's be safe)
        string largeData = string.Join("", Enumerable.Repeat("ABCDEF0123456789", 100));
        var message = new HandshakeRequest { NodeId = largeData, AuthToken = "token" };

        // Act
        await _handler.SendMessageAsync(stream, MessageType.HandshakeReq, message, true, null);

        stream.Position = 0;
        (var type, byte[] payload) = await _handler.ReadMessageAsync(stream, null);

        // Assert
        type.ShouldBe(MessageType.HandshakeReq);
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);
        decoded.NodeId.ShouldBe(largeData);
    }

    /// <summary>
    ///     Verifies an encrypted message can be written and read successfully.
    /// </summary>
    [Fact]
    public async Task RoundTrip_ShouldWork_WithEncryption()
    {
        // Arrange
        var stream = new MemoryStream();
        var message = new HandshakeRequest { NodeId = "secure-node", AuthToken = "secure-token" };

        // Mock CipherState
        var key = new byte[32]; // 256-bit key
        new Random().NextBytes(key);
        var cipherState = new CipherState(key, key); // Encrypt and Decrypt with same key for loopback

        // Act
        await _handler.SendMessageAsync(stream, MessageType.HandshakeReq, message, false, cipherState);

        stream.Position = 0;
        (var type, byte[] payload) = await _handler.ReadMessageAsync(stream, cipherState);

        // Assert
        type.ShouldBe(MessageType.HandshakeReq);
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);
        decoded.NodeId.ShouldBe("secure-node");
    }

    /// <summary>
    ///     Verifies a message can be round-tripped when both compression and encryption are enabled.
    /// </summary>
    [Fact]
    public async Task RoundTrip_ShouldWork_WithEncryption_And_Compression()
    {
        // Arrange
        var stream = new MemoryStream();
        string largeData = string.Join("", Enumerable.Repeat("SECURECOMPRESSION", 100));
        var message = new HandshakeRequest { NodeId = largeData };

        var key = new byte[32];
        new Random().NextBytes(key);
        var cipherState = new CipherState(key, key);

        // Act: Compress THEN Encrypt
        await _handler.SendMessageAsync(stream, MessageType.HandshakeReq, message, true, cipherState);

        stream.Position = 0;
        // Verify wire encryption (should be MessageType.SecureEnv)
        // But ReadMessageAsync abstracts this away. 
        // We can peek at the stream if we want, but let's trust ReadMessageAsync handles it.

        (var type, byte[] payload) = await _handler.ReadMessageAsync(stream, cipherState);

        // Assert
        type.ShouldBe(MessageType.HandshakeReq);
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);
        decoded.NodeId.ShouldBe(largeData);
    }

    /// <summary>
    ///     Verifies that message reads succeed when bytes arrive in small fragments.
    /// </summary>
    [Fact]
    public async Task ReadMessage_ShouldHandle_Fragmentation()
    {
        // Arrange
        var fullStream = new MemoryStream();
        var message = new HandshakeRequest { NodeId = "fragmented" };
        await _handler.SendMessageAsync(fullStream, MessageType.HandshakeReq, message, false, null);

        byte[] completeBytes = fullStream.ToArray();
        var fragmentedStream = new FragmentedMemoryStream(completeBytes, 2); // Read 2 bytes at a time

        // Act
        (var type, byte[] payload) = await _handler.ReadMessageAsync(fragmentedStream, null);

        // Assert
        type.ShouldBe(MessageType.HandshakeReq);
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);
        decoded.NodeId.ShouldBe("fragmented");
    }

    /// <summary>
    ///     Verifies that dataset-aware protocol fields are serialized and parsed correctly.
    /// </summary>
    [Fact]
    public void DatasetAwareMessages_ShouldRoundTripDatasetFields()
    {
        var request = new PullChangesRequest
        {
            SinceWall = 10,
            SinceLogic = 1,
            SinceNode = "node-a",
            DatasetId = "logs"
        };

        byte[] payload = request.ToByteArray();
        var decoded = PullChangesRequest.Parser.ParseFrom(payload);

        decoded.DatasetId.ShouldBe("logs");
    }

    /// <summary>
    ///     Verifies that legacy messages with no dataset id default to the primary dataset.
    /// </summary>
    [Fact]
    public void DatasetAwareMessages_WhenMissingDataset_ShouldDefaultToPrimary()
    {
        var legacy = new HandshakeRequest
        {
            NodeId = "node-legacy",
            AuthToken = "token"
        };

        byte[] payload = legacy.ToByteArray();
        var decoded = HandshakeRequest.Parser.ParseFrom(payload);

        DatasetId.Normalize(decoded.DatasetId).ShouldBe(DatasetId.Primary);
    }

    // Helper Stream for fragmentation test
    private class FragmentedMemoryStream : MemoryStream
    {
        private readonly int _chunkSize;

        /// <summary>
        ///     Initializes a new instance of the <see cref="FragmentedMemoryStream" /> class.
        /// </summary>
        /// <param name="buffer">The backing stream buffer.</param>
        /// <param name="chunkSize">The maximum bytes returned per read.</param>
        public FragmentedMemoryStream(byte[] buffer, int chunkSize) : base(buffer)
        {
            _chunkSize = chunkSize;
        }

        /// <inheritdoc />
        public override async Task<int> ReadAsync(byte[] buffer, int offset, int count,
            CancellationToken cancellationToken)
        {
            // Force read to be max _chunkSize, even if more is requested
            int toRead = Math.Min(count, _chunkSize);
            return await base.ReadAsync(buffer, offset, toRead, cancellationToken);
        }
    }
}