jellyfin-ha-src/tests/Jellyfin.Server.Implementations.Tests/ScheduledTasks/ScanLeaderLeaseTests.cs

using System;
using System.Threading;
using System.Threading.Tasks;
using Emby.Server.Implementations.ScheduledTasks;
using MediaBrowser.Controller.ScheduledTasks;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Moq;
using StackExchange.Redis;
using Xunit;

namespace Jellyfin.Server.Implementations.Tests.ScheduledTasks;

/// <summary>
/// Tests for scan-leader lease behavior. The acquire/renew/takeover state machine is exercised
/// through an in-memory reference implementation that mirrors the Redis Lua contract (no real Redis
/// required), while the fail-safe and success paths of <see cref="RedisScanLeaderLease"/> are
/// exercised against a mocked <see cref="IConnectionMultiplexer"/>.
/// </summary>
public class ScanLeaderLeaseTests
{
    /// <summary>
    /// Verifies that the first instance to call the lease becomes the leader.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task Acquire_WhenUnheld_ReturnsTrue()
    {
        var store = new FakeLeaderStore();
        var clock = new TestClock(DateTime.UtcNow);
        var podA = new ReferenceScanLeaderLease(store, "pod-a", clock, TimeSpan.FromSeconds(60));

        Assert.True(await podA.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
        Assert.Equal("pod-a", store.Owner);
    }

    /// <summary>
    /// Verifies that the current leader renewing its own lease succeeds.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task Renew_BySameInstance_ReturnsTrue()
    {
        var store = new FakeLeaderStore();
        var clock = new TestClock(DateTime.UtcNow);
        var podA = new ReferenceScanLeaderLease(store, "pod-a", clock, TimeSpan.FromSeconds(60));

        Assert.True(await podA.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));

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

        Assert.True(await podA.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
        Assert.Equal("pod-a", store.Owner);
    }

    /// <summary>
    /// Verifies that a second instance cannot acquire the lease while the leader's lease is still valid.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task Acquire_BySecondInstance_WhileLeaseValid_ReturnsFalse()
    {
        var store = new FakeLeaderStore();
        var clock = new TestClock(DateTime.UtcNow);
        var podA = new ReferenceScanLeaderLease(store, "pod-a", clock, TimeSpan.FromSeconds(60));
        var podB = new ReferenceScanLeaderLease(store, "pod-b", clock, TimeSpan.FromSeconds(60));

        Assert.True(await podA.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));

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

        Assert.False(await podB.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
        Assert.Equal("pod-a", store.Owner);
    }

    /// <summary>
    /// Verifies that a second instance takes over the lease once the previous leader's lease has expired.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task Acquire_BySecondInstance_AfterLeaseExpires_ReturnsTrue()
    {
        var store = new FakeLeaderStore();
        var clock = new TestClock(DateTime.UtcNow);
        var podA = new ReferenceScanLeaderLease(store, "pod-a", clock, TimeSpan.FromSeconds(60));
        var podB = new ReferenceScanLeaderLease(store, "pod-b", clock, TimeSpan.FromSeconds(60));

        Assert.True(await podA.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));

        // Advance past pod-a's lease expiry without pod-a renewing.
        clock.Advance(TimeSpan.FromSeconds(61));

        Assert.True(await podB.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
        Assert.Equal("pod-b", store.Owner);
    }

    /// <summary>
    /// Verifies that <see cref="NullScanLeaderLease"/> always reports the caller as the leader.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task NullScanLeaderLease_AlwaysReturnsTrue()
    {
        var lease = new NullScanLeaderLease();

        Assert.True(await lease.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
        Assert.True(await lease.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
    }

    /// <summary>
    /// Verifies that <see cref="RedisScanLeaderLease"/> returns <c>true</c> (fail-safe) when the Redis
    /// evaluation throws, so that scheduled scans keep running when Redis is unreachable.
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task RedisScanLeaderLease_WhenRedisThrows_ReturnsTrue()
    {
        var dbMock = new Mock<IDatabase>();
        dbMock
            .Setup(d => d.ScriptEvaluateAsync(
                It.IsAny<string>(),
                It.IsAny<RedisKey[]>(),
                It.IsAny<RedisValue[]>(),
                It.IsAny<CommandFlags>()))
            .ThrowsAsync(new InvalidOperationException("Redis unavailable"));

        var lease = new RedisScanLeaderLease(
            CreateMultiplexer(dbMock.Object),
            Options.Create(new ScanLeaderOptions { Enabled = true, LeaseDurationSeconds = 60 }),
            new Mock<ILogger<RedisScanLeaderLease>>().Object);

        Assert.True(await lease.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
    }

    /// <summary>
    /// Verifies that <see cref="RedisScanLeaderLease"/> reports leadership when the Redis script
    /// returns 1 (lease acquired or renewed).
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task RedisScanLeaderLease_WhenScriptReturnsOne_ReturnsTrue()
    {
        var dbMock = new Mock<IDatabase>();
        dbMock
            .Setup(d => d.ScriptEvaluateAsync(
                It.IsAny<string>(),
                It.IsAny<RedisKey[]>(),
                It.IsAny<RedisValue[]>(),
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(RedisResult.Create((RedisValue)1L));

        var lease = new RedisScanLeaderLease(
            CreateMultiplexer(dbMock.Object),
            Options.Create(new ScanLeaderOptions { Enabled = true, LeaseDurationSeconds = 60 }),
            new Mock<ILogger<RedisScanLeaderLease>>().Object);

        Assert.True(await lease.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
    }

    /// <summary>
    /// Verifies that <see cref="RedisScanLeaderLease"/> reports non-leadership when the Redis script
    /// returns 0 (another instance holds a live lease).
    /// </summary>
    [Fact]
    [Trait("Category", "UnitTest")]
    public async Task RedisScanLeaderLease_WhenScriptReturnsZero_ReturnsFalse()
    {
        var dbMock = new Mock<IDatabase>();
        dbMock
            .Setup(d => d.ScriptEvaluateAsync(
                It.IsAny<string>(),
                It.IsAny<RedisKey[]>(),
                It.IsAny<RedisValue[]>(),
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(RedisResult.Create((RedisValue)0L));

        var lease = new RedisScanLeaderLease(
            CreateMultiplexer(dbMock.Object),
            Options.Create(new ScanLeaderOptions { Enabled = true, LeaseDurationSeconds = 60 }),
            new Mock<ILogger<RedisScanLeaderLease>>().Object);

        Assert.False(await lease.TryAcquireOrRenewAsync(TestContext.Current.CancellationToken));
    }

    private static IConnectionMultiplexer CreateMultiplexer(IDatabase database)
    {
        var muxMock = new Mock<IConnectionMultiplexer>();
        muxMock
            .Setup(m => m.GetDatabase(It.IsAny<int>(), It.IsAny<object>()))
            .Returns(database);
        return muxMock.Object;
    }

    private sealed class FakeLeaderStore
    {
        public string? Owner { get; set; }

        public DateTime ExpiresUtc { get; set; }
    }

    private sealed class TestClock
    {
        private DateTime _now;

        public TestClock(DateTime now)
        {
            _now = now;
        }

        public DateTime UtcNow => _now;

        public void Advance(TimeSpan by) => _now += by;
    }

    /// <summary>
    /// In-memory reference lease that mirrors the Redis Lua acquire-or-renew contract: a key that is
    /// unset or expired is claimed by the caller; a key already owned by the caller is renewed; a key
    /// owned by a different, still-valid holder is refused.
    /// </summary>
    private sealed class ReferenceScanLeaderLease : IScanLeaderLease
    {
        private readonly FakeLeaderStore _store;
        private readonly string _podId;
        private readonly TestClock _clock;
        private readonly TimeSpan _ttl;

        public ReferenceScanLeaderLease(FakeLeaderStore store, string podId, TestClock clock, TimeSpan ttl)
        {
            _store = store;
            _podId = podId;
            _clock = clock;
            _ttl = ttl;
        }

        public Task<bool> TryAcquireOrRenewAsync(CancellationToken cancellationToken = default)
        {
            var now = _clock.UtcNow;
            var currentOwner = _store.Owner is not null && now < _store.ExpiresUtc ? _store.Owner : null;

            if (currentOwner is null)
            {
                _store.Owner = _podId;
                _store.ExpiresUtc = now + _ttl;
                return Task.FromResult(true);
            }

            if (string.Equals(currentOwner, _podId, StringComparison.Ordinal))
            {
                _store.ExpiresUtc = now + _ttl;
                return Task.FromResult(true);
            }

            return Task.FromResult(false);
        }
    }
}