jellyfin-plugin-enigma2/Tests/AdaptiveProbeTests.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Jellyfin.Plugin.Enigma2;

public static class AdaptiveProbeTests
{
    private const int Sid = 15718;
    private const int PacketSize = 188;
    private static string Metadata(int sid, int width)
    {
        var video = new { index = 0, codec_type = "video", codec_name = "h264", width, height = width == 0 ? 0 : 1080 };
        return JsonSerializer.Serialize(new { streams = new[] { video }, programs = new[] { new { program_id = sid, streams = new[] { video } } }, tags = new { secret = "DO_NOT_EXPORT_SOURCE_TAGS" } });
    }

    public static async Task Run(Action<bool, string> check)
    {
        var limits = new LiveStreamProbe.SamplingLimits(TimeSpan.FromMilliseconds(100), TimeSpan.FromMilliseconds(300), PacketSize * 2, PacketSize * 24);
        var snapshots = new List<string>();
        var attempts = 0;
        var lateStream = new ProbePacketStream(gateAfter: 2, delay: 3);
        using (var handler = new ProbeHttpHandler(lateStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, async (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                check(programme == Sid, "adaptive capture preserves expected programme");
                if (!extended)
                {
                    check(new FileInfo(path).Length == limits.InitialBytes, "first snapshot keeps the fast byte limit");
                    lateStream.Continue();
                    await lateStream.ReachedLateData.Task.WaitAsync(token);
                    check(lateStream.Packets >= 8, "reader continues during first metadata analysis");
                    return Metadata(Sid, 0);
                }
                var data = await File.ReadAllBytesAsync(path, token);
                check(data.Length <= limits.MaximumBytes && data.Contains((byte)0x55), "second snapshot contains late video parameters within byte limit");
                return Metadata(Sid, 1920);
            }, limits, TimeSpan.FromSeconds(2));
            var results = await Task.WhenAll(Enumerable.Range(0, 4).Select(_ => probe.Get("http://relay/channel", CancellationToken.None, Sid)));
            check(attempts == 2 && handler.Calls == 1 && results.All(x => x[0].Width == 1920), "late video succeeds with two analyses and one shared HTTP connection");
            check(lateStream.Disposed, "successful adaptive probe closes its HTTP subscription");
        }
        check(Deleted(snapshots), "successful adaptive probe deletes capture and snapshots");

        snapshots.Clear(); attempts = 0;
        var rejectedSampleStream = new ProbePacketStream(delay: 1);
        using (var handler = new ProbeHttpHandler(rejectedSampleStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                if (!extended) throw new LiveStreamProbe.ProbeProcessException(1);
                return Task.FromResult(Metadata(Sid, 1920));
            }, limits, TimeSpan.FromSeconds(2));
            check((await probe.Get("http://relay/channel", CancellationToken.None, Sid))[0].Width == 1920 && attempts == 2 && handler.Calls == 1,
                "a first ffprobe sample rejection retries analysis on the same HTTP capture");
        }
        check(rejectedSampleStream.Disposed && Deleted(snapshots), "retried process rejection cleans up subscription and sample files");

        snapshots.Clear(); attempts = 0;
        var failedAnalysisStream = new ProbePacketStream(delay: 1);
        using (var handler = new ProbeHttpHandler(failedAnalysisStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                throw new InvalidOperationException("Analysis implementation failed.");
            }, limits, TimeSpan.FromSeconds(2));
            var error = await Failure(probe.Get("http://relay/channel", CancellationToken.None, Sid));
            check(attempts == 1 && handler.Calls == 1 && error.Contains("cause=probe_analysis_failure"),
                "an unrelated first analysis failure is not retried");
        }
        check(failedAnalysisStream.Disposed && Deleted(snapshots), "unrelated analysis failure cleans up subscription and sample files");

        snapshots.Clear(); attempts = 0;
        var initiallyQuietStream = new ProbePacketStream(delay: 150);
        using (var handler = new ProbeHttpHandler(initiallyQuietStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                check(extended && new FileInfo(path).Length > 0, "empty initial window waits for bounded extended capture");
                return Task.FromResult(Metadata(Sid, 1920));
            }, limits, TimeSpan.FromSeconds(2));
            check((await probe.Get("http://relay/channel", CancellationToken.None, Sid))[0].Width == 1920 && attempts == 1 && handler.Calls == 1, "video arriving after an empty first window still succeeds");
        }
        check(initiallyQuietStream.Disposed && Deleted(snapshots), "initially quiet capture deletes all sample files");

        snapshots.Clear(); attempts = 0;
        var wrongStream = new ProbePacketStream(delay: 1);
        using (var handler = new ProbeHttpHandler(wrongStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                return Task.FromResult(Metadata(Sid + 1, 1920));
            }, limits, TimeSpan.FromSeconds(2));
            var error = await Failure(probe.Get("http://relay/channel", CancellationToken.None, Sid));
            check(handler.Calls == 1 && attempts == 2 && error.Contains("cause=missing_expected_programme") && error.Contains("programmes=[15719]") && !error.Contains("DO_NOT_EXPORT_SOURCE_TAGS") && !error.Contains("http://"), "wrong programme remains rejected with safe diagnostic after retry");
        }
        check(wrongStream.Disposed && Deleted(snapshots), "wrong programme failure cleans up subscription and sample files");

        snapshots.Clear(); attempts = 0;
        var limitedStream = new ProbePacketStream(bulk: true);
        var byteLimits = limits with { MaximumBytes = PacketSize * 8 };
        using (var handler = new ProbeHttpHandler(limitedStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                snapshots.Add(path); attempts++;
                check(new FileInfo(path).Length <= (extended ? byteLimits.MaximumBytes : byteLimits.InitialBytes), "each adaptive snapshot has an exact byte ceiling");
                return Task.FromResult(Metadata(Sid, 0));
            }, byteLimits, TimeSpan.FromSeconds(2));
            var error = await Failure(probe.Get("http://relay/channel", CancellationToken.None, Sid));
            check(handler.Calls == 1 && attempts == 2 && limitedStream.BytesRead == byteLimits.MaximumBytes && error.Contains("capturedBytes=1504") && error.Contains("captureStop=byte_limit"), "continuous input stops exactly at maximum size and performs at most two analyses");
        }
        check(limitedStream.Disposed && Deleted(snapshots), "bounded size failure releases all capture resources");

        snapshots.Clear(); attempts = 0;
        var stalledStream = new ProbePacketStream(stalled: true);
        var beforeSilentCapture = Directory.GetFiles(Path.GetTempPath(), "enigma-probe-*.capture.ts*").ToHashSet();
        using (var handler = new ProbeHttpHandler(stalledStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, (path, programme, extended, token) =>
            {
                attempts++; return Task.FromResult(Metadata(Sid, 1920));
            }, limits, TimeSpan.FromMilliseconds(180));
            var error = await Failure(probe.Get("http://relay/channel", CancellationToken.None, Sid));
            check(handler.Calls == 1 && attempts == 0 && stalledStream.Disposed && error.Contains("cause=deadline")
                && !Directory.GetFiles(Path.GetTempPath(), "enigma-probe-*.capture.ts*").Except(beforeSilentCapture).Any(), "silent source obeys total deadline, closes its sole HTTP connection and deletes empty capture");
        }

        snapshots.Clear();
        var cancelledAnalysis = false;
        var analysingStream = new ProbePacketStream(gateAfter: 2);
        using (var handler = new ProbeHttpHandler(analysingStream))
        using (var client = new HttpClient(handler) { Timeout = Timeout.InfiniteTimeSpan })
        {
            var probe = new LiveStreamProbe(client, async (path, programme, extended, token) =>
            {
                snapshots.Add(path);
                try { await Task.Delay(Timeout.Infinite, token); return Metadata(Sid, 1920); }
                finally { cancelledAnalysis = true; }
            }, limits, TimeSpan.FromMilliseconds(180));
            var error = await Failure(probe.Get("http://relay/channel", CancellationToken.None, Sid));
            check(cancelledAnalysis && handler.Calls == 1 && snapshots.Count == 1 && analysingStream.Disposed && Deleted(snapshots) && error.Contains("cause=deadline"), "deadline cancels analysis without retry, joins the reader and removes temporary files");
        }
    }

    private static bool Deleted(IEnumerable<string> paths) => paths.All(path => !File.Exists(path) && !File.Exists(path[..^10]));
    private static async Task<string> Failure(Task task)
    {
        try { await task; throw new Exception("Expected probe failure."); }
        catch (InvalidOperationException error) { return error.InnerException?.Message ?? error.Message; }
    }
}

public sealed class ProbeHttpHandler : HttpMessageHandler
{
    private readonly Stream stream;
    public int Calls;
    public ProbeHttpHandler(Stream stream) { this.stream = stream; }
    protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken token)
    {
        Interlocked.Increment(ref Calls);
        return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK) { Content = new StreamContent(stream) });
    }
}

public sealed class ProbePacketStream : Stream
{
    private readonly int gateAfter;
    private readonly int delay;
    private readonly bool stalled;
    private readonly bool bulk;
    private readonly TaskCompletionSource continued = new(TaskCreationOptions.RunContinuationsAsynchronously);
    public TaskCompletionSource ReachedLateData { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
    public int Packets { get; private set; }
    public long BytesRead { get; private set; }
    public bool Disposed { get; private set; }
    public ProbePacketStream(int gateAfter = 0, int delay = 0, bool stalled = false, bool bulk = false)
    { this.gateAfter = gateAfter; this.delay = delay; this.stalled = stalled; this.bulk = bulk; }
    public void Continue() => continued.TrySetResult();
    public override async ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken token = default)
    {
        if (stalled) await Task.Delay(Timeout.Infinite, token);
        if (gateAfter > 0 && Packets >= gateAfter) await continued.Task.WaitAsync(token);
        if (delay > 0) await Task.Delay(delay, token);
        token.ThrowIfCancellationRequested();
        var count = Math.Min(buffer.Length, bulk ? 188 * 256 : 188);
        buffer.Span[..count].Clear();
        for (var i = 0; i < count; i += 188) { buffer.Span[i] = 0x47; if (i + 3 < count && Packets >= 7) buffer.Span[i + 3] = 0x55; Packets++; }
        BytesRead += count;
        if (Packets >= 8) ReachedLateData.TrySetResult();
        return count;
    }
    protected override void Dispose(bool disposing) { Disposed = true; base.Dispose(disposing); }
    public override bool CanRead => true;
    public override bool CanSeek => false;
    public override bool CanWrite => false;
    public override long Length => throw new NotSupportedException();
    public override long Position { get => 0; set => throw new NotSupportedException(); }
    public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
    public override void Flush() { }
    public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
    public override void SetLength(long value) => throw new NotSupportedException();
    public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
}