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

using AvParser.Core.Proxies;

namespace AvParser.Core.Tests.Proxies;

public class ProxyPoolTests
{
    private static readonly ProxyOptions PoolMode = new()
    {
        HealthCheck = ProxyHealthCheck.Pool,
        Rotation = ProxyRotation.Sticky,
    };

    private static ProxyPool Build(
        out FakeProxySource source,
        out FakeProxyProbe probe,
        out FakeTimeProvider clock,
        ProxyOptions? options = null,
        params string[] hosts
    )
    {
        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 ?? PoolMode, clock);
    }

    [Fact]
    public async Task Refresh_loads_every_source()
    {
        var pool = Build(out _, out _, out _, hosts: ["a", "b", "c"]);

        var count = await pool.RefreshAsync(TestContext.Current.CancellationToken);

        count.ShouldBe(3);
        pool.Entries.Count.ShouldBe(3);
    }

    [Fact]
    public async Task Refresh_drops_duplicates_across_sources()
    {
        var feed = new FakeProxySource(ProxySourceKind.Feed);
        feed.Endpoints.Add(ProxyFactory.Endpoint("shared"));

        var custom = new FakeProxySource(ProxySourceKind.Custom) { Id = "custom" };
        custom.Endpoints.Add(ProxyFactory.Endpoint("shared"));

        var pool = new ProxyPool([feed, custom], new FakeProxyProbe(), PoolMode, new FakeTimeProvider());

        (await pool.RefreshAsync(TestContext.Current.CancellationToken)).ShouldBe(1);
    }

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

        pool.Entries[0].RecordSuccess(clock.GetUtcNow(), TimeSpan.FromMilliseconds(120));

        // Free lists are republished every few minutes; a reload that reset every counter would
        // throw away the only real evidence the app has.
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries[0].SuccessCount.ShouldBe(1);
        pool.Entries[0].Latency!.Value.TotalMilliseconds.ShouldBe(120);
    }

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

        source.Endpoints.RemoveAll(endpoint => endpoint.Host == "b");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

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

    [Fact]
    public async Task Refresh_applies_the_protocol_filter()
    {
        var source = new FakeProxySource(ProxySourceKind.Feed);
        source.Endpoints.Add(ProxyFactory.Endpoint("http", protocol: ProxyProtocol.Http));
        source.Endpoints.Add(ProxyFactory.Endpoint("socks", protocol: ProxyProtocol.Socks5));

        var options = PoolMode with { Protocols = ProxyProtocolFilter.Socks5 };
        var pool = new ProxyPool([source], new FakeProxyProbe(), options, new FakeTimeProvider());

        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries.ShouldHaveSingleItem().Endpoint.Protocol.ShouldBe(ProxyProtocol.Socks5);
    }

    [Fact]
    public async Task Refresh_skips_the_feed_when_it_is_switched_off()
    {
        var pool = Build(out var source, out _, out _, PoolMode with { UseFeed = false }, "a");

        (await pool.RefreshAsync(TestContext.Current.CancellationToken)).ShouldBe(0);
        source.GetCallCount.ShouldBe(0);
    }

    [Fact]
    public async Task Acquire_returns_null_for_an_empty_pool() =>
        (await Build(out _, out _, out _).AcquireAsync(TestContext.Current.CancellationToken)).ShouldBeNull();

    [Fact]
    public async Task Pool_mode_hands_out_without_probing()
    {
        var pool = Build(out _, out var probe, out _, hosts: ["a"]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);

        lease.ShouldNotBeNull();
        probe.ProbeCount.ShouldBe(0);
    }

    [Fact]
    public async Task Lazy_mode_skips_past_dead_proxies()
    {
        var options = PoolMode with { HealthCheck = ProxyHealthCheck.Lazy, Rotation = ProxyRotation.RoundRobin };
        var pool = Build(out _, out var probe, out _, options, "dead1", "dead2", "alive");

        probe.Set(ProxyFactory.Endpoint("alive"), alive: true);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);

        lease.ShouldNotBeNull();
        lease.Endpoint.Host.ShouldBe("alive");
    }

    [Fact]
    public async Task Lazy_mode_gives_up_after_the_configured_number_of_attempts()
    {
        var options = PoolMode with
        {
            HealthCheck = ProxyHealthCheck.Lazy,
            Rotation = ProxyRotation.RoundRobin,
            LazyProbeAttempts = 2,
        };
        var pool = Build(out _, out var probe, out _, options, "a", "b", "c", "d", "e");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        (await pool.AcquireAsync(TestContext.Current.CancellationToken)).ShouldBeNull();
        probe.ProbeCount.ShouldBe(2);
    }

    [Fact]
    public async Task Lazy_mode_trusts_a_proxy_already_known_to_be_alive()
    {
        var options = PoolMode with { HealthCheck = ProxyHealthCheck.Lazy };
        var pool = Build(out _, out var probe, out var clock, options, "a");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries[0].RecordProbe(clock.GetUtcNow(), alive: true, TimeSpan.FromMilliseconds(10), null);

        await pool.AcquireAsync(TestContext.Current.CancellationToken);

        probe.ProbeCount.ShouldBe(0);
    }

    [Fact]
    public async Task A_failing_proxy_is_quarantined_and_comes_back_later()
    {
        var options = PoolMode with
        {
            FailuresBeforeQuarantine = 1,
            BaseQuarantine = TimeSpan.FromSeconds(30),
            MaxQuarantine = TimeSpan.FromMinutes(15),
        };
        var pool = Build(out _, out _, out var clock, options, "a");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        lease.ShouldNotBeNull();
        lease.ReportFailure("boom");

        // Sidelined immediately...
        (await pool.AcquireAsync(TestContext.Current.CancellationToken)).ShouldBeNull();

        clock.Advance(TimeSpan.FromSeconds(31));

        // ...and available again once the window expires.
        (await pool.AcquireAsync(TestContext.Current.CancellationToken)).ShouldNotBeNull();
    }

    [Fact]
    public async Task The_quarantine_window_grows_with_repeated_failures()
    {
        var options = PoolMode with
        {
            FailuresBeforeQuarantine = 1,
            BaseQuarantine = TimeSpan.FromSeconds(10),
            MaxQuarantine = TimeSpan.FromHours(1),
        };
        var pool = Build(out _, out _, out var clock, options, "a");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        var entry = pool.Entries[0];

        entry.RecordFailure(clock.GetUtcNow(), options.BaseQuarantine, options.MaxQuarantine, 1);
        var first = entry.QuarantinedUntilUtc!.Value - clock.GetUtcNow();

        entry.RecordFailure(clock.GetUtcNow(), options.BaseQuarantine, options.MaxQuarantine, 1);
        var second = entry.QuarantinedUntilUtc!.Value - clock.GetUtcNow();

        second.ShouldBeGreaterThan(first);
    }

    [Fact]
    public async Task The_quarantine_window_is_capped()
    {
        var pool = Build(out _, out _, out var clock, hosts: ["a"]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);
        var entry = pool.Entries[0];

        for (var i = 0; i < 40; i++)
        {
            entry.RecordFailure(clock.GetUtcNow(), TimeSpan.FromSeconds(10), TimeSpan.FromMinutes(5), 1);
        }

        (entry.QuarantinedUntilUtc!.Value - clock.GetUtcNow()).ShouldBeLessThanOrEqualTo(TimeSpan.FromMinutes(5));
    }

    [Fact]
    public async Task A_success_clears_the_quarantine()
    {
        var pool = Build(out _, out _, out var clock, PoolMode with { FailuresBeforeQuarantine = 1 }, "a");
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        var first = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        first!.ReportFailure();
        clock.Advance(TimeSpan.FromMinutes(1));

        var second = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        second!.ReportSuccess(TimeSpan.FromMilliseconds(50));

        pool.Entries[0].QuarantinedUntilUtc.ShouldBeNull();
        pool.Entries[0].ConsecutiveFailures.ShouldBe(0);
    }

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

        using (await pool.AcquireAsync(TestContext.Current.CancellationToken))
        {
            // A cancelled operation is not the proxy's fault.
        }

        pool.Entries[0].FailureCount.ShouldBe(0);
        pool.Entries[0].SuccessCount.ShouldBe(0);
    }

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

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        lease!.ReportSuccess();
        lease.ReportFailure("late");

        pool.Entries[0].SuccessCount.ShouldBe(1);
        pool.Entries[0].FailureCount.ShouldBe(0);
    }

    [Fact]
    public async Task Sweep_probes_everything_and_counts_the_survivors()
    {
        var pool = Build(out _, out var probe, out _, hosts: ["a", "b", "c"]);
        probe.Set(ProxyFactory.Endpoint("a"), alive: true);
        probe.Set(ProxyFactory.Endpoint("c"), alive: true);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        var reports = new List<ProxySweepProgress>();
        var alive = await pool.SweepAsync(
            new SynchronousProgress<ProxySweepProgress>(reports.Add),
            TestContext.Current.CancellationToken
        );

        alive.ShouldBe(2);
        probe.ProbeCount.ShouldBe(3);
        reports.Count.ShouldBe(3);
        reports[^1].Fraction.ShouldBe(1d);
    }

    [Fact]
    public async Task Sweep_on_an_empty_pool_reports_completion_rather_than_hanging()
    {
        var pool = Build(out _, out _, out _);
        var reports = new List<ProxySweepProgress>();

        var alive = await pool.SweepAsync(
            new SynchronousProgress<ProxySweepProgress>(reports.Add),
            TestContext.Current.CancellationToken
        );

        alive.ShouldBe(0);
        reports.ShouldHaveSingleItem().Total.ShouldBe(0);
    }

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

        var sticky = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        var stickyAgain = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        stickyAgain!.Endpoint.Key.ShouldBe(sticky!.Endpoint.Key);

        pool.Configure(PoolMode with { Rotation = ProxyRotation.RoundRobin });

        var first = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        var second = await pool.AcquireAsync(TestContext.Current.CancellationToken);
        second!.Endpoint.Key.ShouldNotBe(first!.Endpoint.Key);
    }

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

        var slow = pool.Entries.Single(entry => entry.Endpoint.Host == "slow");
        var fast = pool.Entries.Single(entry => entry.Endpoint.Host == "fast");
        slow.RecordProbe(clock.GetUtcNow(), alive: true, TimeSpan.FromMilliseconds(900), null);
        fast.RecordProbe(clock.GetUtcNow(), alive: true, TimeSpan.FromMilliseconds(30), null);

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);

        lease!.Endpoint.Host.ShouldBe("fast");
    }

    [Fact]
    public async Task Only_confirmed_proxies_are_handed_out_while_any_are_confirmed()
    {
        // Ordering was not enough: round-robin and weighted-random draw from the whole list, so on a
        // feed of a few thousand unchecked addresses nearly every pick was a stranger.
        var pool = Build(
            out _,
            out _,
            out var clock,
            new ProxyOptions { Rotation = ProxyRotation.RoundRobin },
            hosts: ["unchecked-a", "known-good", "unchecked-b", "unchecked-c"]
        );
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.Entries.Single(entry => entry.Endpoint.Host == "known-good")
            .RecordProbe(clock.GetUtcNow(), alive: true, TimeSpan.FromMilliseconds(30), null);

        for (var attempt = 0; attempt < 8; attempt++)
        {
            var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);
            lease.ShouldNotBeNull().Endpoint.Host.ShouldBe("known-good");
        }
    }

    [Fact]
    public async Task With_nothing_confirmed_the_pool_still_hands_something_out()
    {
        // Unknown is not dead. A strict filter on a cold pool would stop the collector rather than
        // let it try, and a proxy that answers on this path marks itself live by succeeding.
        var pool = Build(out _, out _, out _, hosts: ["a", "b"]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        pool.LiveCount.ShouldBe(0);
        (await pool.AcquireAsync(TestContext.Current.CancellationToken)).ShouldNotBeNull();
    }

    [Fact]
    public async Task A_proxy_that_has_just_failed_is_not_preferred_over_a_confirmed_one()
    {
        var pool = Build(out _, out _, out var clock, hosts: ["failed", "known-good"]);
        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        // Dead but not yet quarantined: it used to stay in the draw at full weight.
        pool.Entries.Single(entry => entry.Endpoint.Host == "failed")
            .RecordProbe(clock.GetUtcNow(), alive: false, null, "no route");
        pool.Entries.Single(entry => entry.Endpoint.Host == "known-good")
            .RecordProbe(clock.GetUtcNow(), alive: true, TimeSpan.FromMilliseconds(50), null);

        var lease = await pool.AcquireAsync(TestContext.Current.CancellationToken);

        lease.ShouldNotBeNull().Endpoint.Host.ShouldBe("known-good");
    }

    [Fact]
    public async Task Changed_fires_when_the_pool_moves()
    {
        var pool = Build(out _, out _, out _, hosts: ["a"]);
        var fired = 0;
        pool.Changed += (_, _) => fired++;

        await pool.RefreshAsync(TestContext.Current.CancellationToken);

        fired.ShouldBeGreaterThan(0);
    }

    [Fact]
    public void Invalid_options_are_rejected_rather_than_misbehaving_quietly()
    {
        Should.Throw<ArgumentOutOfRangeException>(() => new ProxyOptions { ProbeConcurrency = 0 }.Validated());
        Should.Throw<ArgumentOutOfRangeException>(() => new ProxyOptions { LazyProbeAttempts = 0 }.Validated());
        Should.Throw<ArgumentOutOfRangeException>(() => new ProxyOptions { ProbeTimeout = TimeSpan.Zero }.Validated());
        Should.Throw<ArgumentOutOfRangeException>(() =>
            new ProxyOptions
            {
                BaseQuarantine = TimeSpan.FromHours(2),
                MaxQuarantine = TimeSpan.FromMinutes(1),
            }.Validated()
        );
    }
}