CBDD/tests/CBDD.Tests/BsonDocumentAndBufferWriterTests.cs

using System.Buffers;
using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Text;
using ZB.MOM.WW.CBDD.Bson;

namespace ZB.MOM.WW.CBDD.Tests;

public class BsonDocumentAndBufferWriterTests
{
    [Fact]
    public void BsonDocument_Create_And_TryGet_RoundTrip()
    {
        var keyMap = new ConcurrentDictionary<string, ushort>(StringComparer.OrdinalIgnoreCase);
        var reverseMap = new ConcurrentDictionary<ushort, string>();

        RegisterKey(keyMap, reverseMap, 1, "name");
        RegisterKey(keyMap, reverseMap, 2, "age");
        RegisterKey(keyMap, reverseMap, 3, "_id");

        var expectedId = ObjectId.NewObjectId();

        var doc = BsonDocument.Create(keyMap, b =>
        {
            b.AddString("name", "Alice");
            b.AddInt32("age", 32);
            b.AddObjectId("_id", expectedId);
        });

        var wrapped = new BsonDocument(doc.RawData.ToArray(), reverseMap);

        wrapped.TryGetString("name", out var name).ShouldBeTrue();
        name.ShouldBe("Alice");

        wrapped.TryGetInt32("age", out var age).ShouldBeTrue();
        age.ShouldBe(32);

        wrapped.TryGetObjectId("_id", out var id).ShouldBeTrue();
        id.ShouldBe(expectedId);

        var reader = wrapped.GetReader();
        reader.ReadDocumentSize().ShouldBeGreaterThan(0);
    }

    [Fact]
    public void BsonDocument_TryGet_Should_Return_False_For_Missing_Or_Wrong_Type()
    {
        var keyMap = new ConcurrentDictionary<string, ushort>(StringComparer.OrdinalIgnoreCase);
        var reverseMap = new ConcurrentDictionary<ushort, string>();

        RegisterKey(keyMap, reverseMap, 1, "name");
        RegisterKey(keyMap, reverseMap, 2, "age");

        var doc = BsonDocument.Create(keyMap, b =>
        {
            b.AddString("name", "Bob");
            b.AddInt32("age", 28);
        });

        var wrapped = new BsonDocument(doc.RawData.ToArray(), reverseMap);

        wrapped.TryGetInt32("name", out _).ShouldBeFalse();
        wrapped.TryGetString("missing", out _).ShouldBeFalse();
        wrapped.TryGetObjectId("age", out _).ShouldBeFalse();
    }

    [Fact]
    public void BsonDocumentBuilder_Should_Grow_Buffer_When_Document_Is_Large()
    {
        var keyMap = new ConcurrentDictionary<string, ushort>(StringComparer.OrdinalIgnoreCase);
        var reverseMap = new ConcurrentDictionary<ushort, string>();

        for (ushort i = 1; i <= 180; i++)
        {
            var key = $"k{i}";
            RegisterKey(keyMap, reverseMap, i, key);
        }

        var builder = new BsonDocumentBuilder(keyMap);
        for (int i = 1; i <= 180; i++)
        {
            builder.AddInt32($"k{i}", i);
        }

        var doc = builder.Build();
        doc.Size.ShouldBeGreaterThan(1024);

        var wrapped = new BsonDocument(doc.RawData.ToArray(), reverseMap);
        wrapped.TryGetInt32("k180", out var value).ShouldBeTrue();
        value.ShouldBe(180);
    }

    [Fact]
    public void BsonBufferWriter_Should_Write_Nested_Document_And_Array()
    {
        var output = new ArrayBufferWriter<byte>();
        var writer = new BsonBufferWriter(output);

        int rootSizePos = writer.BeginDocument();

        int childSizePos = writer.BeginDocument("child");
        writer.WriteString("name", "nested");
        writer.WriteBoolean("active", true);
        writer.EndDocument(childSizePos);
        int childEnd = writer.Position;

        int arraySizePos = writer.BeginArray("nums");
        writer.WriteInt32("0", 1);
        writer.WriteInt32("1", 2);
        writer.EndArray(arraySizePos);
        int arrayEnd = writer.Position;

        writer.EndDocument(rootSizePos);
        int rootEnd = writer.Position;

        var bytes = output.WrittenSpan.ToArray();
        PatchDocumentSize(bytes, childSizePos, childEnd);
        PatchDocumentSize(bytes, arraySizePos, arrayEnd);
        PatchDocumentSize(bytes, rootSizePos, rootEnd);

        var reader = new BsonSpanReader(bytes, new ConcurrentDictionary<ushort, string>());

        reader.ReadDocumentSize().ShouldBe(bytes.Length);

        reader.ReadBsonType().ShouldBe(BsonType.Document);
        reader.ReadCString().ShouldBe("child");
        reader.ReadDocumentSize().ShouldBeGreaterThan(8);

        reader.ReadBsonType().ShouldBe(BsonType.String);
        reader.ReadCString().ShouldBe("name");
        reader.ReadString().ShouldBe("nested");

        reader.ReadBsonType().ShouldBe(BsonType.Boolean);
        reader.ReadCString().ShouldBe("active");
        reader.ReadBoolean().ShouldBeTrue();
        reader.ReadBsonType().ShouldBe(BsonType.EndOfDocument);

        reader.ReadBsonType().ShouldBe(BsonType.Array);
        reader.ReadCString().ShouldBe("nums");
        reader.ReadDocumentSize().ShouldBeGreaterThan(8);

        reader.ReadBsonType().ShouldBe(BsonType.Int32);
        reader.ReadCString().ShouldBe("0");
        reader.ReadInt32().ShouldBe(1);

        reader.ReadBsonType().ShouldBe(BsonType.Int32);
        reader.ReadCString().ShouldBe("1");
        reader.ReadInt32().ShouldBe(2);
        reader.ReadBsonType().ShouldBe(BsonType.EndOfDocument);

        reader.ReadBsonType().ShouldBe(BsonType.EndOfDocument);
    }

    [Fact]
    public void BsonSpanReader_ReadByte_And_ReadCStringSpan_Should_Work()
    {
        var singleByteReader = new BsonSpanReader(new byte[] { 0x2A }, new ConcurrentDictionary<ushort, string>());
        singleByteReader.ReadByte().ShouldBe((byte)0x2A);

        var cstring = Encoding.UTF8.GetBytes("hello\0");
        var cstringReader = new BsonSpanReader(cstring, new ConcurrentDictionary<ushort, string>());
        var destination = new char[16];
        var written = cstringReader.ReadCString(destination);

        new string(destination, 0, written).ShouldBe("hello");
    }

    private static void RegisterKey(
        ConcurrentDictionary<string, ushort> keyMap,
        ConcurrentDictionary<ushort, string> reverseMap,
        ushort id,
        string key)
    {
        keyMap[key] = id;
        reverseMap[id] = key;
    }

    private static void PatchDocumentSize(byte[] output, int sizePosition, int endPosition)
    {
        BinaryPrimitives.WriteInt32LittleEndian(output.AsSpan(sizePosition, 4), endPosition - sizePosition);
    }
}