ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.Communication.Tests/Grpc/SiteStreamEventBatcherTests.cs

using System.Diagnostics;
using System.Threading.Channels;
using Google.Protobuf.WellKnownTypes;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;

namespace ZB.MOM.WW.ScadaBridge.Communication.Tests.Grpc;

/// <summary>
/// Unit tests for the per-subscriber coalescing pump behind R2 (gRPC event batching).
/// The pump is the only new behaviour on the site→central hot path, so its contract is
/// pinned directly rather than only through the server: the size cap, the time cap, the
/// flush when the channel writer completes, the "never reorder" guarantee, and the
/// un-negotiated (maxBatchEvents == 1) shape that keeps an older central working.
/// </summary>
public class SiteStreamEventBatcherTests
{
    private const string Corr = "corr-batch";

    private static SiteStreamEvent Event(int seq) => new()
    {
        CorrelationId = Corr,
        AttributeChanged = new AttributeValueUpdate
        {
            InstanceUniqueName = "SiteA.Pump01",
            AttributePath = "Modules.IO",
            AttributeName = "Seq",
            Value = seq.ToString(),
            Quality = Quality.Good,
            Timestamp = Timestamp.FromDateTimeOffset(DateTimeOffset.UnixEpoch.AddSeconds(seq))
        }
    };

    /// <summary>Flattens an emitted frame into the sequence numbers it carried, in order.</summary>
    private static IEnumerable<int> Seqs(SiteStreamEvent frame)
    {
        if (frame.EventCase == SiteStreamEvent.EventOneofCase.Batch)
        {
            foreach (var inner in frame.Batch.Events)
                yield return int.Parse(inner.AttributeChanged.Value);
            yield break;
        }

        yield return int.Parse(frame.AttributeChanged.Value);
    }

    private sealed record PumpRun(List<SiteStreamEvent> Frames, List<int> FrameSizes, Task Completion);

    private static PumpRun StartPump(
        ChannelReader<SiteStreamEvent> reader,
        int maxBatchEvents,
        TimeSpan window,
        CancellationToken ct = default)
    {
        var frames = new List<SiteStreamEvent>();
        var sizes = new List<int>();
        var task = SiteStreamEventBatcher.PumpAsync(
            reader,
            Corr,
            maxBatchEvents,
            window,
            (evt, _) =>
            {
                lock (frames) { frames.Add(evt); }
                return Task.CompletedTask;
            },
            size => { lock (frames) { sizes.Add(size); } },
            ct);
        return new PumpRun(frames, sizes, task);
    }

    // ── Size cap ────────────────────────────────────────────────────────────────

    [Fact]
    public async Task SizeCap_SplitsABacklogIntoFramesOfAtMostMaxEvents()
    {
        // A backlog already sitting in the channel is drained without waiting, but never
        // beyond the size cap — 250 queued events at a cap of 100 must come out as
        // 100 + 100 + 50, in order, with nothing lost or duplicated.
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        for (var i = 0; i < 250; i++)
            Assert.True(channel.Writer.TryWrite(Event(i)));
        channel.Writer.Complete();

        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromMilliseconds(25));
        await run.Completion;

        Assert.All(run.FrameSizes, s => Assert.True(s <= 100, $"frame carried {s} events (cap 100)"));
        Assert.Equal([100, 100, 50], run.FrameSizes);
        Assert.Equal(Enumerable.Range(0, 250), run.Frames.SelectMany(Seqs));
    }

    [Fact]
    public async Task SizeCapOfOne_EmitsPlainPerEventFrames_TheUnnegotiatedShape()
    {
        // maxBatchEvents == 1 is what an un-negotiated subscription (an older central)
        // gets. Every event must ride its own plain frame — never a Batch case, which
        // that central's generated code cannot parse.
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        for (var i = 0; i < 5; i++)
            channel.Writer.TryWrite(Event(i));
        channel.Writer.Complete();

        var run = StartPump(channel.Reader, maxBatchEvents: 1, window: TimeSpan.FromMilliseconds(25));
        await run.Completion;

        Assert.Equal(5, run.Frames.Count);
        Assert.All(run.Frames, f =>
            Assert.Equal(SiteStreamEvent.EventOneofCase.AttributeChanged, f.EventCase));
        Assert.All(run.Frames, f => Assert.Equal(Corr, f.CorrelationId));
        Assert.Equal(Enumerable.Range(0, 5), run.Frames.SelectMany(Seqs));
    }

    // ── Time cap ────────────────────────────────────────────────────────────────

    [Fact]
    public async Task TimeCap_ClosesAnUnderfullBatchWhenTheWindowElapses()
    {
        // Two events arrive (a backlog, so the pump lingers), then the source goes quiet
        // well short of the size cap. The window — not the cap — must close the batch,
        // and it must do so within a bounded time rather than waiting for a 100th event
        // that never comes.
        var window = TimeSpan.FromMilliseconds(120);
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        channel.Writer.TryWrite(Event(0));
        channel.Writer.TryWrite(Event(1));

        var started = Stopwatch.GetTimestamp();
        var run = StartPump(channel.Reader, maxBatchEvents: 100, window);

        SiteStreamEvent frame;
        while (true)
        {
            lock (run.Frames)
            {
                if (run.Frames.Count > 0) { frame = run.Frames[0]; break; }
            }
            Assert.True(Stopwatch.GetElapsedTime(started) < TimeSpan.FromSeconds(5),
                "the window never closed the underfull batch");
            await Task.Delay(5);
        }

        var elapsed = Stopwatch.GetElapsedTime(started);

        channel.Writer.Complete();
        await run.Completion;

        Assert.Equal(SiteStreamEvent.EventOneofCase.Batch, frame.EventCase);
        Assert.Equal([0, 1], Seqs(frame));
        // The batch waited (it did not close instantly on the two queued events) and it
        // closed on the window, not on a cap it never reached.
        Assert.True(elapsed >= window - TimeSpan.FromMilliseconds(20),
            $"batch closed after {elapsed.TotalMilliseconds:0.0} ms, before the {window.TotalMilliseconds:0} ms window");
    }

    [Fact]
    public async Task LoneEventOnAQuietStream_IsNeverDelayedByTheWindow()
    {
        // The latency contract: the window applies only AFTER a backlog has been observed.
        // A single event on an idle stream must be emitted immediately as a plain frame,
        // so per-event latency on a quiet site is unchanged by batching.
        var window = TimeSpan.FromSeconds(5);
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        var run = StartPump(channel.Reader, maxBatchEvents: 100, window);

        var started = Stopwatch.GetTimestamp();
        channel.Writer.TryWrite(Event(7));

        while (true)
        {
            lock (run.Frames)
            {
                if (run.Frames.Count > 0) break;
            }
            Assert.True(Stopwatch.GetElapsedTime(started) < TimeSpan.FromSeconds(3),
                "a lone event was held by the coalescing window");
            await Task.Delay(2);
        }

        var elapsed = Stopwatch.GetElapsedTime(started);
        channel.Writer.Complete();
        await run.Completion;

        Assert.Equal(SiteStreamEvent.EventOneofCase.AttributeChanged, run.Frames[0].EventCase);
        Assert.True(elapsed < TimeSpan.FromSeconds(1),
            $"lone event took {elapsed.TotalMilliseconds:0.0} ms against a {window.TotalSeconds:0} s window");
    }

    // ── Flush on stream close ───────────────────────────────────────────────────

    [Fact]
    public async Task WriterCompletion_FlushesTheInFlightBatchBeforeReturning()
    {
        // The channel writer completing mid-window (the site stopping the relay actor and
        // calling channel.Writer.TryComplete()) must flush what is already buffered rather
        // than silently discarding it while waiting out the window.
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        channel.Writer.TryWrite(Event(0));
        channel.Writer.TryWrite(Event(1));

        // A long window guarantees the pump is lingering, not already past the emit.
        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromSeconds(30));

        await Task.Delay(100);
        lock (run.Frames)
        {
            Assert.Empty(run.Frames); // still lingering
        }

        channel.Writer.Complete();
        await run.Completion.WaitAsync(TimeSpan.FromSeconds(5));

        Assert.Single(run.Frames);
        Assert.Equal([0, 1], Seqs(run.Frames[0]));
    }

    [Fact]
    public async Task WriterCompletionWithNothingBuffered_ReturnsWithoutEmitting()
    {
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromMilliseconds(25));

        channel.Writer.Complete();
        await run.Completion.WaitAsync(TimeSpan.FromSeconds(5));

        Assert.Empty(run.Frames);
    }

    [Fact]
    public async Task Cancellation_EndsThePumpWithOperationCanceled()
    {
        // Client disconnect / duplicate replacement / site shutdown. The pump must
        // surface OperationCanceledException exactly as the pre-batching await-foreach
        // loop did, so SiteStreamGrpcServer's existing catch and finally are unchanged.
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        using var cts = new CancellationTokenSource();
        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromMilliseconds(25), cts.Token);

        await cts.CancelAsync();

        await Assert.ThrowsAnyAsync<OperationCanceledException>(
            () => run.Completion.WaitAsync(TimeSpan.FromSeconds(5)));
    }

    // ── Ordering ────────────────────────────────────────────────────────────────

    [Fact]
    public async Task Ordering_IsPreservedAcrossManyBatchesUnderAProducerRace()
    {
        // Batching is a framing change and nothing else: with a producer writing
        // concurrently with the pump, the flattened output must be the exact input
        // sequence — no reordering, no loss, no duplication, across many frames.
        const int total = 5_000;
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        var run = StartPump(channel.Reader, maxBatchEvents: 32, window: TimeSpan.FromMilliseconds(5));

        var producer = Task.Run(async () =>
        {
            for (var i = 0; i < total; i++)
            {
                channel.Writer.TryWrite(Event(i));
                if (i % 250 == 0) await Task.Yield();
            }
            channel.Writer.Complete();
        });

        await producer;
        await run.Completion.WaitAsync(TimeSpan.FromSeconds(30));

        Assert.Equal(Enumerable.Range(0, total), run.Frames.SelectMany(Seqs));
        Assert.All(run.FrameSizes, s => Assert.InRange(s, 1, 32));
        Assert.Equal(total, run.FrameSizes.Sum());
    }

    // ── Frame shape ─────────────────────────────────────────────────────────────

    [Fact]
    public async Task BatchFrame_CarriesTheCorrelationIdOnceAndBlanksItOnInnerEvents()
    {
        // The byte saving batching exists for: the correlation id is stamped once on the
        // enclosing frame, not repeated on every packed event. No consumer reads the
        // inner value (SiteStreamGrpcClient.ForEachEvent ignores it).
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        for (var i = 0; i < 4; i++) channel.Writer.TryWrite(Event(i));
        channel.Writer.Complete();

        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromMilliseconds(25));
        await run.Completion;

        var frame = Assert.Single(run.Frames);
        Assert.Equal(SiteStreamEvent.EventOneofCase.Batch, frame.EventCase);
        Assert.Equal(Corr, frame.CorrelationId);
        Assert.All(frame.Batch.Events, e => Assert.Equal(string.Empty, e.CorrelationId));
    }

    [Fact]
    public async Task PerEventTimestampsSurviveBatching()
    {
        // End-to-end latency measurement rides the per-event Timestamp; coalescing must
        // not rewrite it to a single frame-level stamp.
        var channel = Channel.CreateUnbounded<SiteStreamEvent>();
        for (var i = 0; i < 3; i++) channel.Writer.TryWrite(Event(i));
        channel.Writer.Complete();

        var run = StartPump(channel.Reader, maxBatchEvents: 100, window: TimeSpan.FromMilliseconds(25));
        await run.Completion;

        var frame = Assert.Single(run.Frames);
        Assert.Equal(
            [
                DateTimeOffset.UnixEpoch,
                DateTimeOffset.UnixEpoch.AddSeconds(1),
                DateTimeOffset.UnixEpoch.AddSeconds(2)
            ],
            frame.Batch.Events.Select(e => e.AttributeChanged.Timestamp.ToDateTimeOffset()));
    }
}