natsnet/dotnet/tests/ZB.MOM.NatsNet.Server.IntegrationTests/Helpers/CheckHelper.cs

// Copyright 2012-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Mirrors Go checkFor from server/test_test.go.

using System.Diagnostics;
using ZB.MOM.NatsNet.Server;

namespace ZB.MOM.NatsNet.Server.IntegrationTests.Helpers;

/// <summary>
/// Retry/polling helpers for integration tests.
/// Mirrors Go <c>checkFor</c> from server/test_test.go.
/// </summary>
internal static class CheckHelper
{
    /// <summary>
    /// Polls <paramref name="check"/> repeatedly until it returns null (success)
    /// or the timeout expires, in which case the last exception is thrown.
    /// Mirrors Go <c>checkFor(t, timeout, interval, func() error)</c>.
    /// </summary>
    public static void CheckFor(TimeSpan timeout, TimeSpan interval, Func<Exception?> check)
    {
        var sw = Stopwatch.StartNew();
        Exception? last = null;
        while (sw.Elapsed < timeout)
        {
            last = check();
            if (last == null) return;
            Thread.Sleep(interval);
        }

        // One final attempt after the sleep boundary.
        last = check();
        if (last == null) return;

        throw new TimeoutException(
            $"CheckFor timed out after {timeout}: {last.Message}", last);
    }

    /// <summary>
    /// Async version of <see cref="CheckFor"/>. Uses <c>Task.Delay</c> instead of
    /// <c>Thread.Sleep</c> to avoid blocking the thread pool.
    /// </summary>
    public static async Task CheckForAsync(
        TimeSpan timeout,
        TimeSpan interval,
        Func<Task<Exception?>> check,
        CancellationToken cancellationToken = default)
    {
        var sw = Stopwatch.StartNew();
        Exception? last = null;
        while (sw.Elapsed < timeout)
        {
            last = await check().ConfigureAwait(false);
            if (last == null) return;
            await Task.Delay(interval, cancellationToken).ConfigureAwait(false);
        }

        // One final attempt.
        last = await check().ConfigureAwait(false);
        if (last == null) return;

        throw new TimeoutException(
            $"CheckForAsync timed out after {timeout}: {last.Message}", last);
    }

    /// <summary>
    /// Waits until all servers in <paramref name="servers"/> have formed a cluster
    /// (each server sees at least <c>servers.Length - 1</c> routes).
    /// Uses a 10-second timeout with 100 ms poll interval.
    /// Mirrors Go <c>checkClusterFormed</c>.
    /// </summary>
    public static void CheckClusterFormed(params NatsServer[] servers)
    {
        var expected = servers.Length - 1;
        CheckFor(TimeSpan.FromSeconds(10), TimeSpan.FromMilliseconds(100), () =>
        {
            foreach (var s in servers)
            {
                var routes = s.NumRoutes();
                if (routes < expected)
                    return new Exception(
                        $"Server {s.Options.ServerName} has {routes} routes, expected {expected}.");
            }
            return null;
        });
    }

    /// <summary>
    /// Waits until the given server has at least <paramref name="expected"/>
    /// leaf node connections.
    /// Uses a 10-second timeout with 100 ms poll interval.
    /// Mirrors Go <c>checkLeafNodeConnectedCount</c>.
    /// </summary>
    public static void CheckLeafNodeConnectedCount(NatsServer server, int expected)
    {
        CheckFor(TimeSpan.FromSeconds(10), TimeSpan.FromMilliseconds(100), () =>
        {
            var count = server.NumLeafNodes();
            if (count < expected)
                return new Exception(
                    $"Server {server.Options.ServerName} has {count} leaf nodes, expected {expected}.");
            return null;
        });
    }
}