jellyfin-ha-src/tests/Jellyfin.MediaEncoding.Tests/Transcoding/InMemoryTranscodeSessionStoreTests.cs

using System;
using System.Threading.Tasks;
using Jellyfin.MediaEncoding.Tests.Fakes;
using MediaBrowser.Controller.MediaEncoding;
using Xunit;

namespace Jellyfin.MediaEncoding.Tests.Transcoding;

/// <summary>
/// Unit tests for the <see cref="ITranscodeSessionStore"/> contract — lease expiry, double-claim
/// prevention, ownership-checked renewal and stale-session cleanup — against
/// <see cref="InMemoryTranscodeSessionStore"/>, the reference implementation. The Redis-backed
/// implementation is covered by <c>RedisTranscodeSessionStoreTests</c> against a real Redis.
/// </summary>
public class InMemoryTranscodeSessionStoreTests
{
    private static TranscodeSession CreateSession(
        string id,
        string pod,
        DateTime leaseExpiry,
        int lastSegmentIndex = 0,
        long lastOffset = 0L)
        => new TranscodeSession
        {
            PlaySessionId = id,
            OwnerPod = pod,
            LeaseExpiresUtc = leaseExpiry,
            ManifestPath = $"/transcode/{id}/manifest.m3u8",
            SegmentPathPrefix = $"/transcode/{id}/segment",
            MediaSourceId = $"media-source-{id}",
            LastCompletedSegmentIndex = lastSegmentIndex,
            LastDurablePlaybackOffset = lastOffset,
        };

    /// <summary>
    /// Lease expiry: <see cref="ITranscodeSessionStore.TryGetAsync"/> returns <c>null</c>
    /// once <see cref="TranscodeSession.LeaseExpiresUtc"/> has passed.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task TryGetAsync_AfterLeaseExpires_ReturnsNull()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-expired", "pod-a", DateTime.UtcNow.AddMilliseconds(-1));
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        var result = await store.TryGetAsync("session-expired", TestContext.Current.CancellationToken);

        Assert.Null(result);
    }

    /// <summary>
    /// A session whose lease has not yet expired is returned correctly.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task TryGetAsync_WithinLease_ReturnsSession()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-live", "pod-a", DateTime.UtcNow.AddMinutes(5), lastSegmentIndex: 3, lastOffset: 18_000_000L);
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        var result = await store.TryGetAsync("session-live", TestContext.Current.CancellationToken);

        Assert.NotNull(result);
        Assert.Equal("session-live", result.PlaySessionId);
        Assert.Equal("pod-a", result.OwnerPod);
        Assert.Equal(3, result.LastCompletedSegmentIndex);
        Assert.Equal(18_000_000L, result.LastDurablePlaybackOffset);
    }

    /// <summary>
    /// Double-claim prevention: <see cref="ITranscodeSessionStore.TryTakeoverAsync"/> returns
    /// <c>false</c> while the first pod's lease is still valid.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task TryTakeoverAsync_WhileLeaseValid_ReturnsFalse()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-valid", "pod-a", DateTime.UtcNow.AddMinutes(5));
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        var firstAttempt = await store.TryTakeoverAsync("session-valid", "pod-b", TestContext.Current.CancellationToken);
        var secondAttempt = await store.TryTakeoverAsync("session-valid", "pod-c", TestContext.Current.CancellationToken);

        Assert.False(firstAttempt);
        Assert.False(secondAttempt);
    }

    /// <summary>
    /// After a lease expires, the first concurrent caller that invokes
    /// <see cref="ITranscodeSessionStore.TryTakeoverAsync"/> wins; the second caller returns <c>false</c>.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task TryTakeoverAsync_AfterLeaseExpires_OnlyFirstClaimerSucceeds()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-stale", "pod-a", DateTime.UtcNow.AddMilliseconds(-1));
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        // First pod wins; its takeover renews the lease atomically.
        var firstTakeover = await store.TryTakeoverAsync("session-stale", "pod-b", TestContext.Current.CancellationToken);

        // Second pod is too late – pod-b already holds a fresh lease.
        var secondTakeover = await store.TryTakeoverAsync("session-stale", "pod-c", TestContext.Current.CancellationToken);

        Assert.True(firstTakeover);
        Assert.False(secondTakeover);
    }

    /// <summary>
    /// Heartbeat renewal: <see cref="ITranscodeSessionStore.RenewLeaseAsync"/> extends
    /// <see cref="TranscodeSession.LeaseExpiresUtc"/> beyond its original value.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task RenewLeaseAsync_ExtendsLeaseExpiry()
    {
        var store = new InMemoryTranscodeSessionStore();
        var originalExpiry = DateTime.UtcNow.AddSeconds(5);
        var session = CreateSession("session-renew", "pod-a", originalExpiry, lastSegmentIndex: 2, lastOffset: 10_000_000L);
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        Assert.True(await store.RenewLeaseAsync("session-renew", "pod-a", TestContext.Current.CancellationToken));

        var renewed = await store.TryGetAsync("session-renew", TestContext.Current.CancellationToken);
        Assert.NotNull(renewed);
        Assert.True(
            renewed.LeaseExpiresUtc > originalExpiry,
            "Renewed lease expiry should be later than the original expiry.");
    }

    /// <summary>
    /// A renewal from a pod that no longer owns the lease must fail and must not revert ownership.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task RenewLeaseAsync_ByNonOwner_ReturnsFalse()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-renew-other", "pod-a", DateTime.UtcNow.AddMilliseconds(-1));
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        Assert.True(await store.TryTakeoverAsync("session-renew-other", "pod-b", TestContext.Current.CancellationToken));
        Assert.False(await store.RenewLeaseAsync("session-renew-other", "pod-a", TestContext.Current.CancellationToken));

        var current = await store.TryGetAsync("session-renew-other", TestContext.Current.CancellationToken);
        Assert.NotNull(current);
        Assert.Equal("pod-b", current.OwnerPod);
    }

    /// <summary>
    /// Stale-session cleanup: an expired session can be deleted without error, and a
    /// subsequent <see cref="ITranscodeSessionStore.TryGetAsync"/> returns <c>null</c>.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task DeleteAsync_ExpiredSession_CompletesWithoutError()
    {
        var store = new InMemoryTranscodeSessionStore();
        var session = CreateSession("session-delete", "pod-a", DateTime.UtcNow.AddMilliseconds(-1));
        await store.SetAsync(session, TestContext.Current.CancellationToken);

        var ex = await Record.ExceptionAsync(() => store.DeleteAsync("session-delete", TestContext.Current.CancellationToken));
        Assert.Null(ex);

        var result = await store.TryGetAsync("session-delete", TestContext.Current.CancellationToken);
        Assert.Null(result);
    }

    /// <summary>
    /// Deleting a session that was never stored must complete without error.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task DeleteAsync_NonExistentSession_CompletesWithoutError()
    {
        var store = new InMemoryTranscodeSessionStore();

        var ex = await Record.ExceptionAsync(() => store.DeleteAsync("nonexistent-session", TestContext.Current.CancellationToken));

        Assert.Null(ex);
    }
}