av-parser/tests/AvParser.Core.Tests/Proxies/ProxyPoolWarmUpTests.cs

using AvParser.Core.Proxies;

namespace AvParser.Core.Tests.Proxies;

public class ProxyPoolWarmUpTests
{
    private static readonly ProxyOptions Options = new() { ProbeConcurrency = 1 };

    private static ProxyPool Build(out FakeProxyProbe probe, out FakeTimeProvider clock, params string[] hosts)
    {
        var source = new FakeProxySource(ProxySourceKind.Feed);
        source.Endpoints.AddRange(hosts.Select(host => ProxyFactory.Endpoint(host)));

        probe = new FakeProxyProbe();
        clock = new FakeTimeProvider();

        return new ProxyPool([source], probe, Options, clock);
    }

    [Fact]
    public async Task What_worked_last_time_is_tried_first()
    {
        // The whole point of remembering: a second launch should confirm the known-good ones
        // rather than walking a list of a few thousand addresses from the top.
        var pool = Build(out var probe, out var clock, "cold-a", "known-good", "cold-b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries.Single(entry => entry.Endpoint.Host == "known-good")
            .RestoreState(wasAlive: true, TimeSpan.FromMilliseconds(30), 12, 1, clock.GetUtcNow());

        pool.WarmUpOrder()[0].Endpoint.Host.ShouldBe("known-good");

        // Asking for two: one is already known live, so the warm-up has a reason to probe at all.
        probe.DefaultAlive = true;
        await pool.WarmUpAsync(2, cancellationToken: TestContext.Current.CancellationToken);

        probe.Probed[0].ShouldContain("known-good");
    }

    [Fact]
    public async Task The_faster_of_two_remembered_proxies_goes_first()
    {
        var pool = Build(out _, out var clock, "slow", "fast");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries.Single(entry => entry.Endpoint.Host == "slow")
            .RestoreState(wasAlive: true, TimeSpan.FromMilliseconds(900), 3, 0, clock.GetUtcNow());
        pool.Entries.Single(entry => entry.Endpoint.Host == "fast")
            .RestoreState(wasAlive: true, TimeSpan.FromMilliseconds(40), 3, 0, clock.GetUtcNow());

        pool.WarmUpOrder().Select(entry => entry.Endpoint.Host).ShouldBe(["fast", "slow"]);
    }

    [Fact]
    public async Task Warming_up_stops_once_the_target_is_met()
    {
        // Probing all 200 when 2 were asked for would turn every launch into a full sweep.
        var pool = Build(out var probe, out _, [.. Enumerable.Range(0, 200).Select(i => $"h{i}")]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        probe.DefaultAlive = true;

        var live = await pool.WarmUpAsync(2, cancellationToken: TestContext.Current.CancellationToken);

        live.ShouldBeGreaterThanOrEqualTo(2);
        probe.ProbeCount.ShouldBeLessThan(200);
    }

    [Fact]
    public async Task Warming_up_probes_nothing_when_enough_are_already_live()
    {
        var pool = Build(out var probe, out var clock, "a", "b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        foreach (var entry in pool.Entries)
        {
            entry.RecordSuccess(clock.GetUtcNow(), TimeSpan.FromMilliseconds(10));
        }

        (await pool.WarmUpAsync(2, cancellationToken: TestContext.Current.CancellationToken)).ShouldBe(2);
        probe.ProbeCount.ShouldBe(0);
    }

    [Fact]
    public async Task A_quarantined_proxy_is_not_warmed_up()
    {
        var pool = Build(out _, out var clock, "a", "b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries.Single(entry => entry.Endpoint.Host == "a")
            .RecordFailure(clock.GetUtcNow(), TimeSpan.FromMinutes(5), TimeSpan.FromHours(1), 1);

        pool.WarmUpOrder().Select(entry => entry.Endpoint.Host).ShouldBe(["b"]);
    }

    [Fact]
    public async Task Nothing_live_reads_as_nothing_live()
    {
        var pool = Build(out var probe, out _, "a", "b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        probe.DefaultAlive = false;

        (await pool.WarmUpAsync(1, cancellationToken: TestContext.Current.CancellationToken)).ShouldBe(0);
        pool.LiveCount.ShouldBe(0);
    }

    [Fact]
    public async Task The_top_up_checks_what_the_warm_up_skipped()
    {
        // The warm-up leaves almost everything unknown by design; without this pass the pool looks
        // — and behaves — as if it held two proxies rather than the fifty that answer.
        var pool = Build(out var probe, out _, [.. Enumerable.Range(0, 50).Select(index => $"h{index}")]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        probe.DefaultAlive = true;

        await pool.WarmUpAsync(2, cancellationToken: TestContext.Current.CancellationToken);
        var afterWarmUp = probe.ProbeCount;

        afterWarmUp.ShouldBeLessThan(50);

        var found = await pool.TopUpAsync(4, cancellationToken: TestContext.Current.CancellationToken);

        found.ShouldBe(50 - afterWarmUp);
        pool.LiveCount.ShouldBe(50);
        pool.Entries.ShouldAllBe(entry => entry.Health == ProxyHealthState.Alive);
    }

    [Fact]
    public async Task The_top_up_does_not_re_probe_what_is_already_known()
    {
        var pool = Build(out var probe, out _, "a", "b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        probe.DefaultAlive = true;

        await pool.TopUpAsync(2, cancellationToken: TestContext.Current.CancellationToken);
        var first = probe.ProbeCount;

        // Everything has a verdict now, so a second pass has nothing to do — re-probing would just
        // be a sweep, and a sweep is something the user asks for.
        (await pool.TopUpAsync(2, cancellationToken: TestContext.Current.CancellationToken)).ShouldBe(0);
        probe.ProbeCount.ShouldBe(first);
    }

    [Fact]
    public async Task A_stopped_top_up_keeps_what_it_learned()
    {
        var pool = Build(out var probe, out _, [.. Enumerable.Range(0, 20).Select(index => $"h{index}")]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        probe.DefaultAlive = true;

        using var stop = new CancellationTokenSource();
        await stop.CancelAsync();

        await Should.ThrowAsync<OperationCanceledException>(() => pool.TopUpAsync(2, null, stop.Token));

        // Cancelled before anything was probed, so nothing is claimed to be live either.
        pool.LiveCount.ShouldBe(0);
    }
}