/* ========================================================================
 * Copyright (c) 2005-2025 The OPC Foundation, Inc. All rights reserved.
 *
 * OPC Foundation MIT License 1.00
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 *
 * The complete license agreement can be found here:
 * http://opcfoundation.org/License/MIT/1.00/
 * ======================================================================*/

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Reflection;
using System.Threading;
using System.Threading.Tasks;
using Moq;
using NUnit.Framework;
using Opc.Ua.Tests;

namespace Opc.Ua.Client.Tests
{
    [TestFixture]
    [Parallelizable]
    [Category("Client")]
    [Category("ClassicSubscriptionEngine")]
    public class ClassicSubscriptionEngineTests
    {
        private ITelemetryContext m_telemetry;
        private Mock<ISubscriptionEngineContext> m_mockContext;
        private SemaphoreSlim m_lock;

        [SetUp]
        public void SetUp()
        {
            m_telemetry = NUnitTelemetryContext.Create();
            m_lock = new SemaphoreSlim(1, 1);
            m_mockContext = new Mock<ISubscriptionEngineContext>();
            m_mockContext.Setup(c => c.SessionId)
                .Returns(new NodeId(1));
            m_mockContext.Setup(c => c.Disposed)
                .Returns(false);
            m_mockContext.Setup(c => c.Reconnecting)
                .Returns(false);
            m_mockContext.Setup(c => c.Connected)
                .Returns(false);
            m_mockContext.Setup(c => c.Closing)
                .Returns(false);
            m_mockContext.Setup(c => c.Telemetry)
                .Returns(m_telemetry);
            m_mockContext.Setup(c => c.Subscriptions)
                .Returns([]);
            m_mockContext.Setup(c => c.OperationTimeout)
                .Returns(60000);
            m_mockContext.Setup(c => c.ReturnDiagnostics)
                .Returns(DiagnosticsMasks.None);
            m_mockContext.Setup(c => c.ReconnectLock)
                .Returns(m_lock);
            m_mockContext.Setup(c => c.GoodPublishRequestCount)
                .Returns(0);
        }

        [TearDown]
        public void TearDown()
        {
            m_lock.Dispose();
        }

        [Test]
        public void FactoryCreateReturnsClassicEngine()
        {
            ISubscriptionEngine engine =
                ClassicSubscriptionEngineFactory.Instance.Create(
                    m_mockContext.Object);

            Assert.That(engine, Is.Not.Null);
            Assert.That(
                engine,
                Is.InstanceOf<ClassicSubscriptionEngine>());

            engine.Dispose();
        }

        [Test]
        public void PublishResponsePreservesServerSequenceAvailability()
        {
            using var subscription = new Subscription(m_telemetry);
            m_mockContext.Setup(context => context.Subscriptions).Returns([subscription]);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            ArrayOf<uint> available = new uint[] { 9, 10, 11 }.ToArrayOf();

            engine.ProcessPublishResponse(
                new ResponseHeader { Timestamp = DateTimeUtc.Now },
                subscription.Id,
                available,
                false,
                new NotificationMessage
                {
                    SequenceNumber = 11,
                    PublishTime = DateTimeUtc.Now,
                    NotificationData = [new ExtensionObject(new DataChangeNotification())]
                });

            Assert.That(subscription.AvailableSequenceNumbers.ToArray(), Is.EquivalentTo(available.ToArray()!));
        }

        /// <summary>
        /// A Republish response has no available sequence numbers. It must
        /// neither wipe the subscription's list (so other missing messages are
        /// still republished) nor drop pending acknowledgements (L7-1).
        /// </summary>
        [Test]
        public void RepublishResponseKeepsAvailableSequenceNumbersAndAcknowledgements()
        {
            m_mockContext.Setup(c => c.ServerState).Returns(ServerState.Running);
            using var subscription = new Subscription(m_telemetry);
            m_mockContext.Setup(context => context.Subscriptions).Returns([subscription]);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            ArrayOf<uint> available = new uint[] { 7, 8, 9 }.ToArrayOf();

            engine.ProcessPublishResponse(
                new ResponseHeader { Timestamp = DateTimeUtc.Now },
                subscription.Id,
                available,
                false,
                new NotificationMessage
                {
                    SequenceNumber = 9,
                    PublishTime = DateTimeUtc.Now,
                    NotificationData = [new ExtensionObject(new DataChangeNotification())]
                });
            engine.AddPendingAcknowledgement(subscription.Id, 30);

            engine.ProcessPublishResponse(
                new ResponseHeader { Timestamp = DateTimeUtc.Now },
                subscription.Id,
                default,
                false,
                new NotificationMessage
                {
                    SequenceNumber = 7,
                    PublishTime = DateTimeUtc.Now,
                    NotificationData = [new ExtensionObject(new DataChangeNotification())]
                },
                republished: true);

            Assert.That(subscription.AvailableSequenceNumbers.ToArray(), Is.EquivalentTo(available.ToArray()!));
            Assert.That(engine.RemoveAcknowledgementsForSubscription(subscription.Id), Is.EqualTo(3),
                "The acks for 9 and 30 are kept and the republished 7 is added.");
        }

        [Test]
        public void StartPublishingWithNoSubscriptionsDoesNothing()
        {
            m_mockContext.Setup(c => c.Subscriptions)
                .Returns([]);

            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            engine.StartPublishing(timeout: 5000, fullQueue: false);

            m_mockContext.Verify(
                c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()),
                Times.Never);
        }

        [Test]
        public void NotifySubscriptionsChangedTriggersPublishReEvaluation()
        {
            m_mockContext.Setup(c => c.Subscriptions)
                .Returns([]);
            m_mockContext.Setup(c => c.Disposed)
                .Returns(false);

            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            // With no subscriptions, NotifySubscriptionsChanged
            // should not trigger a publish request.
            engine.NotifySubscriptionsChanged();

            m_mockContext.Verify(
                c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()),
                Times.Never);
        }

        [Test]
        public void MinPublishRequestCountDefaultsCorrectly()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(
                engine.MinPublishRequestCount, Is.EqualTo(1));
            Assert.That(
                engine.MaxPublishRequestCount,
                Is.EqualTo(ushort.MaxValue));
        }

        [Test]
        public void MinMaxPublishRequestCountCanBeSet()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            engine.MinPublishRequestCount = 5;
            engine.MaxPublishRequestCount = 50;

            Assert.That(
                engine.MinPublishRequestCount, Is.EqualTo(5));
            Assert.That(
                engine.MaxPublishRequestCount, Is.EqualTo(50));
        }

        [Test]
        public void MinPublishRequestCountRejectsInvalidValues()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(
                () => engine.MinPublishRequestCount = 0,
                Throws.TypeOf<ArgumentOutOfRangeException>());
            Assert.That(
                () => engine.MinPublishRequestCount = 101,
                Throws.TypeOf<ArgumentOutOfRangeException>());
        }

        [Test]
        public void MaxPublishRequestCountRejectsInvalidValues()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(
                () => engine.MaxPublishRequestCount = 0,
                Throws.TypeOf<ArgumentOutOfRangeException>());
        }

        [Test]
        public void PauseAndResumePublishing()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(engine.PausePublishing, Throws.Nothing);
            Assert.That(engine.ResumePublishing, Throws.Nothing);
        }

        [Test]
        public async Task StopPublishingAsyncDisposesCleanlyAsync()
        {
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);

            await engine.StopPublishingAsync().ConfigureAwait(false);
        }

        [Test]
        public void DisposeIsIdempotent()
        {
#pragma warning disable CA2000 // Dispose objects before losing scope
            var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);
#pragma warning restore CA2000 // Dispose objects before losing scope

            Assert.That(engine.Dispose, Throws.Nothing);
            Assert.That(engine.Dispose, Throws.Nothing);
        }

        [Test]
        public void ConstructorThrowsOnNullContext()
        {
            Assert.That(
                () => new ClassicSubscriptionEngine(null!),
                Throws.TypeOf<ArgumentNullException>());
        }

        [Test]
        public void GoodPublishRequestCountDelegatesToContext()
        {
            m_mockContext.Setup(c => c.GoodPublishRequestCount)
                .Returns(42);

            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(
                engine.GoodPublishRequestCount, Is.EqualTo(42));
        }

        [Test]
        public void BadPublishRequestCountIsZero()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            Assert.That(
                engine.BadPublishRequestCount, Is.Zero);
        }

        [Test]
        public void BeginPublishSkipsWhenSessionKeepAliveStopped()
        {
            // Arrange — session looks healthy except KeepAliveStopped is true.
            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.Reconnecting).Returns(false);
            m_mockContext.Setup(c => c.Closing).Returns(false);
            m_mockContext.Setup(c => c.KeepAliveStopped).Returns(true);
            m_mockContext.Setup(c => c.LastKeepAliveTime)
                .Returns(DateTime.UtcNow.AddSeconds(-30));

            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            // Act
            bool result = engine.BeginPublish(timeout: 5000);

            // Assert — guard short-circuited; no PublishAsync was invoked.
            Assert.That(result, Is.False);
            m_mockContext.Verify(
                c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()),
                Times.Never);
        }

        [Test]
        public void MaxPublishRequestCountRespectsMinimum()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            engine.MinPublishRequestCount = 10;
            engine.MaxPublishRequestCount = 5;

            // MaxPublishRequestCount getter returns the larger of
            // min and max to ensure min is always respected.
            Assert.That(
                engine.MaxPublishRequestCount,
                Is.GreaterThanOrEqualTo(
                    engine.MinPublishRequestCount));
        }

        /// <summary>
        /// Regression coverage for the V1 ack-pruning helper added by
        /// #3540. <see cref="ClassicSubscriptionEngine.RemoveAcknowledgementsForSubscription"/>
        /// drops every queued acknowledgement targeting the given
        /// subscription id while leaving acks for other subscriptions
        /// untouched. The recovery path relies on this before recreating
        /// a subscription so stale acks for the dead id cannot produce
        /// <c>BadSubscriptionIdInvalid</c> responses on servers that
        /// re-use subscription identifiers (e.g. Kepware always
        /// starting at <c>1</c>).
        /// </summary>
        [Test]
        public void RemoveAcknowledgementsForSubscriptionDropsOnlyMatchingAcks()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            engine.AddPendingAcknowledgement(1u, 10u);
            engine.AddPendingAcknowledgement(2u, 11u);
            engine.AddPendingAcknowledgement(1u, 12u);
            engine.AddPendingAcknowledgement(1u, 13u);
            engine.AddPendingAcknowledgement(3u, 14u);

            int dropped = engine.RemoveAcknowledgementsForSubscription(1u);

            Assert.That(dropped, Is.EqualTo(3),
                "All queued acks for the dead subscription id must be dropped.");

            int droppedAgain = engine.RemoveAcknowledgementsForSubscription(1u);
            Assert.That(droppedAgain, Is.Zero,
                "Second call must drop nothing (idempotent).");

            int droppedSecond = engine.RemoveAcknowledgementsForSubscription(2u);
            Assert.That(droppedSecond, Is.EqualTo(1),
                "Survivors must still be in the queue and drainable per-id.");

            int droppedThird = engine.RemoveAcknowledgementsForSubscription(3u);
            Assert.That(droppedThird, Is.EqualTo(1));
        }

        /// <summary>
        /// Calling
        /// <see cref="ClassicSubscriptionEngine.RemoveAcknowledgementsForSubscription"/>
        /// on an empty pending-ack list is a no-op and must not throw.
        /// </summary>
        [Test]
        public void RemoveAcknowledgementsForSubscriptionOnEmptyQueueReturnsZero()
        {
            using var engine =
                new ClassicSubscriptionEngine(m_mockContext.Object);

            int dropped = engine.RemoveAcknowledgementsForSubscription(42u);
            Assert.That(dropped, Is.Zero);
        }

        [Test]
        public void ProcessPublishResponseHandlesKeepAliveUnknownAndAcknowledgements()
        {
            m_mockContext.Setup(c => c.ServerState).Returns(ServerState.Running);
            m_mockContext.Setup(c => c.SessionId).Returns(new NodeId(1));
            m_mockContext.Setup(c => c.Subscriptions).Returns([]);
            m_mockContext.Setup(c => c.DeleteSubscriptionsOnClose).Returns(false);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            engine.AddPendingAcknowledgement(7, 10);
            engine.AddPendingAcknowledgement(7, 200);

            engine.ProcessPublishResponse(
                new ResponseHeader
                {
                    Timestamp = DateTime.UtcNow,
                    StringTable = ["s"]
                },
                7,
                [10u, 200u],
                false,
                new NotificationMessage
                {
                    SequenceNumber = 201,
                    PublishTime = DateTime.UtcNow,
                    NotificationData = []
                });

            m_mockContext.Verify(c => c.OnKeepAlive(ServerState.Running, It.IsAny<DateTime>()), Times.Once);
            m_mockContext.Verify(c => c.OnPublishNotification(
                It.IsAny<Subscription>(),
                It.IsAny<NotificationEventArgs>()), Times.Never);
        }

        /// <summary>
        /// A pending acknowledgement exactly 2^31 away from the latest
        /// sequence number must not throw from Math.Abs(int.MinValue) and
        /// abort the publish response (L4-1).
        /// </summary>
        [Test]
        public void ProcessPublishResponseToleratesAcknowledgementHalfTheSequenceSpaceAway()
        {
            m_mockContext.Setup(c => c.ServerState).Returns(ServerState.Running);
            m_mockContext.Setup(c => c.Subscriptions).Returns([]);
            m_mockContext.Setup(c => c.DeleteSubscriptionsOnClose).Returns(false);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            engine.AddPendingAcknowledgement(7, 0x80000000u);

            Assert.That(() => engine.ProcessPublishResponse(
                new ResponseHeader { Timestamp = DateTime.UtcNow },
                7,
                [],
                false,
                new NotificationMessage
                {
                    SequenceNumber = 1,
                    PublishTime = DateTime.UtcNow,
                    NotificationData = []
                }), Throws.Nothing);

            Assert.That(engine.RemoveAcknowledgementsForSubscription(7), Is.Zero,
                "An acknowledgement far outside the tolerance is dropped.");
        }

        [TestCase(10u, 1u, true)]
        [TestCase(1u, 10u, true)]
        [TestCase(1u, uint.MaxValue, true)]
        [TestCase(20u, 1u, false)]
        [TestCase(0x80000000u, 0u, false)]
        [TestCase(0u, 0x80000000u, false)]
        public void OutOfOrderThresholdIsWrapAwareAndNeverThrows(
            uint sequenceNumber, uint latest, bool expected)
        {
            Assert.That(
                ClassicSubscriptionEngine.IsWithinOutOfOrderThreshold(sequenceNumber, latest),
                Is.EqualTo(expected));
        }

        [Test]
        public async Task ProcessPublishResponseDispatchesKnownSubscriptionNotificationAsync()
        {
            m_mockContext.Setup(c => c.ServerState).Returns(ServerState.Running);
            var session = CreateSubscriptionSession(7);
            using var subscription = new Subscription(
                m_telemetry,
                new SubscriptionOptions
                {
                    PublishingEnabled = false
                })
            {
                Session = session.Object
            };
            await subscription.CreateAsync().ConfigureAwait(false);
            m_mockContext.Setup(c => c.Subscriptions).Returns([subscription]);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            var dataChange = new DataChangeNotification
            {
                MonitoredItems = [],
                DiagnosticInfos = []
            };

            engine.ProcessPublishResponse(
                new ResponseHeader { Timestamp = DateTime.UtcNow, StringTable = ["ns"] },
                subscription.Id,
                [2u],
                true,
                new NotificationMessage
                {
                    SequenceNumber = 2,
                    PublishTime = DateTime.UtcNow,
                    NotificationData = [new ExtensionObject(dataChange)]
                });

            m_mockContext.Verify(c => c.OnPublishNotification(
                subscription,
                It.Is<NotificationEventArgs>(args => args.NotificationMessage.SequenceNumber == 2)),
                Times.Once);
        }

        [Test]
        public void ProcessRepublishResponseErrorClassifiesServiceResults()
        {
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            (uint statusCode, bool expectedHandled)[] cases =
            [
                (StatusCodes.BadSubscriptionIdInvalid.Code, true),
                (StatusCodes.BadMessageNotAvailable.Code, true),
                (StatusCodes.BadEncodingLimitsExceeded.Code, true),
                (StatusCodes.BadUnexpectedError.Code, false)
            ];

            foreach ((uint statusCode, bool expectedHandled) in cases)
            {
                (bool handled, ServiceResult error) = engine.ProcessRepublishResponseError(
                    new ServiceResultException(statusCode),
                    11,
                    22);

                Assert.That(handled, Is.EqualTo(expectedHandled));
                Assert.That(error.StatusCode, Is.EqualTo(statusCode));
            }

            m_mockContext.Verify(c => c.OnPublishError(
                It.IsAny<ServiceResult>(),
                11,
                22), Times.Exactly(cases.Length));
        }

        [Test]
        public async Task OnPublishCompleteProcessesSuccessfulResponseAndRequeuesAsync()
        {
            var sessionId = new NodeId(1);
            m_mockContext.Setup(c => c.SessionId).Returns(sessionId);
            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.ServerState).Returns(ServerState.Running);
            var session = CreateSubscriptionSession(1);
            using var subscription = new Subscription(
                m_telemetry,
                new SubscriptionOptions
                {
                    PublishingEnabled = false
                })
            {
                Session = session.Object
            };
            await subscription.CreateAsync().ConfigureAwait(false);
            m_mockContext.Setup(c => c.Subscriptions).Returns([subscription]);
            m_mockContext.Setup(c => c.PrepareAcknowledgementsToSend(
                    It.IsAny<List<SubscriptionAcknowledgement>>()))
                .Returns<List<SubscriptionAcknowledgement>>(acknowledgements =>
                    (new List<SubscriptionAcknowledgement>(acknowledgements), []));
            var followUpResponse = new TaskCompletionSource<PublishResponse>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            var publishCalled = new TaskCompletionSource<bool>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            var followUpAcknowledgements = ArrayOf<SubscriptionAcknowledgement>.Empty;
            m_mockContext.Setup(c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()))
                .Callback<RequestHeader, ArrayOf<SubscriptionAcknowledgement>, CancellationToken>(
                    (_, acknowledgements, _) =>
                    {
                        followUpAcknowledgements = acknowledgements;
                        publishCalled.TrySetResult(true);
                    })
                .Returns(new ValueTask<PublishResponse>(followUpResponse.Task));
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            Task<PublishResponse> task = Task.FromResult(new PublishResponse
            {
                ResponseHeader = new ResponseHeader
                {
                    Timestamp = DateTime.UtcNow,
                    StringTable = []
                },
                SubscriptionId = 1,
                AvailableSequenceNumbers = [],
                MoreNotifications = false,
                NotificationMessage = new NotificationMessage
                {
                    SequenceNumber = 1,
                    PublishTime = DateTime.UtcNow,
                    NotificationData =
                    [
                        new ExtensionObject(new DataChangeNotification
                        {
                            MonitoredItems = [],
                            DiagnosticInfos = []
                        })
                    ]
                },
                Results = [StatusCodes.Good]
            });
            var sentAcknowledgements = new List<SubscriptionAcknowledgement>
            {
                new() { SubscriptionId = subscription.Id, SequenceNumber = 5 }
            };

            InvokeOnPublishComplete(engine, task, sessionId, sentAcknowledgements);
            Task completed = await Task.WhenAny(
                publishCalled.Task,
                Task.Delay(TimeSpan.FromSeconds(5))).ConfigureAwait(false);

            Assert.That(completed, Is.SameAs(publishCalled.Task));
            Assert.That(followUpAcknowledgements, Has.Count.EqualTo(1));
            Assert.That(followUpAcknowledgements[0].SubscriptionId, Is.EqualTo(subscription.Id));
            Assert.That(followUpAcknowledgements[0].SequenceNumber, Is.EqualTo(1u));
            m_mockContext.Verify(c => c.PrepareAcknowledgementsToSend(
                It.IsAny<List<SubscriptionAcknowledgement>>()), Times.Once);
            m_mockContext.Verify(c => c.PublishAsync(
                It.IsAny<RequestHeader>(),
                It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                It.IsAny<CancellationToken>()), Times.Once);
            m_mockContext.Verify(c => c.AsyncRequestCompleted(
                task,
                99,
                DataTypes.PublishRequest), Times.Once);
            m_mockContext.Verify(c => c.AsyncRequestStarted(
                followUpResponse.Task,
                It.IsAny<System.Diagnostics.Activity>(),
                It.IsAny<uint>(),
                DataTypes.PublishRequest), Times.Once);
            m_mockContext.Verify(c => c.OnKeepAlive(ServerState.Running, It.IsAny<DateTime>()), Times.Once);
        }

        [Test]
        public void OnPublishCompleteRaisesPublishErrorForFaultedResponse()
        {
            m_mockContext.Setup(c => c.SessionId).Returns(new NodeId(1));
            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.Subscriptions).Returns([
                new Subscription(m_telemetry)
                {
                    CurrentPublishingInterval = 100,
                    CurrentLifetimeCount = 10
                }
            ]);
            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);
            Task<PublishResponse> task = Task.FromException<PublishResponse>(
                new ServiceResultException(StatusCodes.BadSessionIdInvalid));

            InvokeOnPublishComplete(
                engine,
                task,
                new NodeId(1),
                [new SubscriptionAcknowledgement { SubscriptionId = 1, SequenceNumber = 2 }]);

            m_mockContext.Verify(c => c.OnPublishError(
                It.Is<ServiceResult>(result => result.StatusCode == StatusCodes.BadSessionIdInvalid),
                0,
                0), Times.Once);
            m_mockContext.Verify(c => c.OnKeepAliveError(
                It.Is<ServiceResult>(result => result.StatusCode == StatusCodes.BadSessionIdInvalid)), Times.Once);
        }

        [Test]
        public void ConcurrentPublishReEvaluationDoesNotExceedDesiredRequestCount()
        {
            const int subscriptionCount = 5;
            const int concurrentCompletions = 64;

            var subscriptions = new List<Subscription>();
            for (int ii = 0; ii < subscriptionCount; ii++)
            {
                subscriptions.Add(
                    new Subscription(m_telemetry)
                    {
                        CurrentPublishingInterval = 100,
                        CurrentLifetimeCount = 10
                    });
            }

            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.Subscriptions).Returns(subscriptions);
            m_mockContext.Setup(c => c.PrepareAcknowledgementsToSend(
                    It.IsAny<List<SubscriptionAcknowledgement>>()))
                .Returns(([], []));

            // Mirrors the session: a publish request only becomes visible to
            // GoodPublishRequestCount once AsyncRequestStarted has recorded it,
            // which the engine does after the request has already been issued.
            int outstanding = 0;
            m_mockContext.Setup(c => c.GoodPublishRequestCount)
                .Returns(() => Volatile.Read(ref outstanding));
            m_mockContext.Setup(c => c.AsyncRequestStarted(
                    It.IsAny<Task>(),
                    It.IsAny<Activity>(),
                    It.IsAny<uint>(),
                    It.IsAny<uint>()))
                .Callback(() => Interlocked.Increment(ref outstanding));

            // Requests never complete, so nothing drains the outstanding count.
            // RunContinuationsAsynchronously keeps the completion below off
            // this thread, so cancelling cannot re-enter the engine inline and
            // issue further publishes while the assertion is evaluated.
            var pending = new TaskCompletionSource<PublishResponse>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            int issued = 0;
            m_mockContext.Setup(c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()))
                .Returns(() =>
                {
                    Interlocked.Increment(ref issued);
                    return new ValueTask<PublishResponse>(pending.Task);
                });

            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);

            using var barrier = new Barrier(concurrentCompletions);
            var threads = new Task[concurrentCompletions];
            for (int ii = 0; ii < concurrentCompletions; ii++)
            {
                threads[ii] = Task.Factory.StartNew(
                    () =>
                    {
                        barrier.SignalAndWait();
                        engine.NotifySubscriptionsChanged();
                    },
                    CancellationToken.None,
                    TaskCreationOptions.LongRunning,
                    TaskScheduler.Default);
            }

            Task.WaitAll(threads);

            // Sample before completing the pending requests so the result is
            // fixed by the re-evaluation alone.
            int issuedByReEvaluation = Volatile.Read(ref issued);
            pending.TrySetCanceled();

            Assert.That(
                issuedByReEvaluation,
                Is.LessThanOrEqualTo(subscriptionCount),
                "Concurrent publish re-evaluation must not issue more requests " +
                "than the desired publish request count.");
        }

        [Test]
        public void StartPublishingRefillsPipelineAfterRequestsAreWrittenOff()
        {
            const int subscriptionCount = 5;

            var subscriptions = new List<Subscription>();
            for (int ii = 0; ii < subscriptionCount; ii++)
            {
                var subscription = new Subscription(m_telemetry)
                {
                    CurrentPublishingInterval = 100,
                    CurrentLifetimeCount = 10
                };

                // Created is Id != 0, and the server assigns the id, so the
                // desired publish count only counts subscriptions the server
                // has acknowledged.
                typeof(Subscription).GetProperty(
                        nameof(Subscription.Id),
                        BindingFlags.Instance | BindingFlags.Public)!
                    .SetValue(subscription, (uint)(ii + 1));

                subscriptions.Add(subscription);
            }

            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.Subscriptions).Returns(subscriptions);
            m_mockContext.Setup(c => c.PrepareAcknowledgementsToSend(
                    It.IsAny<List<SubscriptionAcknowledgement>>()))
                .Returns(([], []));

            // Mirrors Session.GoodPublishRequestCount: a request is recorded by
            // AsyncRequestStarted and stops counting once it is written off.
            int outstanding = 0;
            int writtenOff = 0;
            m_mockContext.Setup(c => c.GoodPublishRequestCount)
                .Returns(() => Volatile.Read(ref outstanding) - Volatile.Read(ref writtenOff));
            m_mockContext.Setup(c => c.AsyncRequestStarted(
                    It.IsAny<Task>(),
                    It.IsAny<Activity>(),
                    It.IsAny<uint>(),
                    It.IsAny<uint>()))
                .Callback(() => Interlocked.Increment(ref outstanding));

            // The requests never complete: this is the state Session.OnKeepAlive
            // recovers from, where the responses are outstanding but are no
            // longer expected to return. RunContinuationsAsynchronously keeps
            // the completion below off the test thread: the engine hangs its
            // continuation on this task with OnCompleted, so an inline
            // completion would re-enter the engine synchronously and issue
            // further publishes while the assertion is being evaluated.
            var pending = new TaskCompletionSource<PublishResponse>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            int issued = 0;
            m_mockContext.Setup(c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()))
                .Returns(() =>
                {
                    Interlocked.Increment(ref issued);
                    return new ValueTask<PublishResponse>(pending.Task);
                });

            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);

            engine.StartPublishing(timeout: 5000, fullQueue: true);
            int afterFill = Volatile.Read(ref issued);
            Assert.That(afterFill, Is.EqualTo(subscriptionCount));

            // Session.OnKeepAlive marks every outstanding publish request
            // defunct and then calls StartPublishing to refill the pipeline.
            Volatile.Write(ref writtenOff, afterFill);
            Assert.That(m_mockContext.Object.GoodPublishRequestCount, Is.Zero);

            engine.StartPublishing(timeout: 5000, fullQueue: false);

            // Sample before completing the pending requests so the result is
            // fixed by StartPublishing alone.
            int afterRefill = Volatile.Read(ref issued);
            pending.TrySetCanceled();

            Assert.That(
                afterRefill - afterFill,
                Is.EqualTo(subscriptionCount),
                "StartPublishing must refill a pipeline whose requests were " +
                "written off, otherwise the session never publishes again.");
        }

        [Test]
        public void BeginPublishIsBoundedByTheDesiredRequestCount()
        {
            const int subscriptionCount = 3;

            var subscriptions = new List<Subscription>();
            for (int ii = 0; ii < subscriptionCount; ii++)
            {
                subscriptions.Add(
                    new Subscription(m_telemetry)
                    {
                        CurrentPublishingInterval = 100,
                        CurrentLifetimeCount = 10
                    });
            }

            m_mockContext.Setup(c => c.Connected).Returns(true);
            m_mockContext.Setup(c => c.Subscriptions).Returns(subscriptions);
            m_mockContext.Setup(c => c.PrepareAcknowledgementsToSend(
                    It.IsAny<List<SubscriptionAcknowledgement>>()))
                .Returns(([], []));

            int outstanding = 0;
            int writtenOff = 0;
            m_mockContext.Setup(c => c.GoodPublishRequestCount)
                .Returns(() => Volatile.Read(ref outstanding) - Volatile.Read(ref writtenOff));
            m_mockContext.Setup(c => c.AsyncRequestStarted(
                    It.IsAny<Task>(),
                    It.IsAny<Activity>(),
                    It.IsAny<uint>(),
                    It.IsAny<uint>()))
                .Callback(() => Interlocked.Increment(ref outstanding));

            var pending = new TaskCompletionSource<PublishResponse>(
                TaskCreationOptions.RunContinuationsAsynchronously);
            int issued = 0;
            m_mockContext.Setup(c => c.PublishAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<ArrayOf<SubscriptionAcknowledgement>>(),
                    It.IsAny<CancellationToken>()))
                .Returns(() =>
                {
                    Interlocked.Increment(ref issued);
                    return new ValueTask<PublishResponse>(pending.Task);
                });

            using var engine = new ClassicSubscriptionEngine(m_mockContext.Object);

            // Every subscription nudges the pipeline when it sees no
            // notification, which is what Subscription.HandleOnKeepAliveStopped
            // does. The nudges must not multiply the pipeline.
            for (int ii = 0; ii < subscriptionCount * 4; ii++)
            {
                engine.BeginPublish(timeout: 5000);
            }

            int afterNudges = Volatile.Read(ref issued);
            Assert.That(
                afterNudges,
                Is.EqualTo(subscriptionCount),
                "The keep alive nudge must not send past the desired publish " +
                "request count.");

            // Once the outstanding requests are written off the pipeline is
            // drained, so the nudge has to get through again.
            Volatile.Write(ref writtenOff, afterNudges);

            Assert.That(engine.BeginPublish(timeout: 5000), Is.True);

            int afterRecovery = Volatile.Read(ref issued);
            pending.TrySetCanceled();

            Assert.That(afterRecovery, Is.EqualTo(afterNudges + 1));
        }

        private static void InvokeOnPublishComplete(
            ClassicSubscriptionEngine engine,
            Task<PublishResponse> task,
            NodeId sessionId,
            List<SubscriptionAcknowledgement>? acknowledgements)
        {
            MethodInfo method = typeof(ClassicSubscriptionEngine).GetMethod(
                "OnPublishComplete",
                BindingFlags.Instance | BindingFlags.NonPublic)!;
            method.Invoke(
                engine,
                [
                    task,
                    sessionId,
                    acknowledgements,
                    new RequestHeader { RequestHandle = 99 }
                ]);
        }

        private static Mock<ISession> CreateSubscriptionSession(uint subscriptionId)
        {
            var session = new Mock<ISession>(MockBehavior.Loose);
            session.SetupGet(s => s.SessionTimeout).Returns(60000);
            session.Setup(s => s.CreateSubscriptionAsync(
                    It.IsAny<RequestHeader>(),
                    It.IsAny<double>(),
                    It.IsAny<uint>(),
                    It.IsAny<uint>(),
                    It.IsAny<uint>(),
                    It.IsAny<bool>(),
                    It.IsAny<byte>(),
                    It.IsAny<CancellationToken>()))
                .Returns(new ValueTask<CreateSubscriptionResponse>(new CreateSubscriptionResponse
                {
                    ResponseHeader = new ResponseHeader(),
                    SubscriptionId = subscriptionId,
                    RevisedPublishingInterval = 100,
                    RevisedMaxKeepAliveCount = 10,
                    RevisedLifetimeCount = 30
                }));
            return session;
        }
    }
}