fengling-gateway/tests/YarpGateway.Tests/Unit/Services/RedisConnectionManagerTests.cs

using Microsoft.Extensions.Logging;
using Moq;
using StackExchange.Redis;
using Xunit;
using FluentAssertions;
using YarpGateway.Config;
using YarpGateway.Services;

namespace YarpGateway.Tests.Unit.Services;

public class RedisConnectionManagerTests
{
    private readonly Mock<IConnectionMultiplexer> _connectionMock;
    private readonly Mock<IDatabase> _databaseMock;
    private readonly Mock<ILogger<RedisConnectionManager>> _loggerMock;
    private readonly RedisConfig _config;

    public RedisConnectionManagerTests()
    {
        _connectionMock = new Mock<IConnectionMultiplexer>();
        _databaseMock = new Mock<IDatabase>();
        
        _connectionMock
            .Setup(x => x.GetDatabase(It.IsAny<int>(), It.IsAny<object>()))
            .Returns(_databaseMock.Object);
        
        _loggerMock = new Mock<ILogger<RedisConnectionManager>>();
        
        _config = new RedisConfig
        {
            ConnectionString = "localhost:6379",
            InstanceName = "test-instance",
            Database = 0
        };
    }

    private RedisConnectionManager CreateManager(IConnectionMultiplexer? connection = null)
    {
        var conn = connection ?? _connectionMock.Object;
        
        // Use reflection to create the manager with a mock connection
        var manager = new RedisConnectionManager(_config, _loggerMock.Object);
        
        // Replace the lazy connection
        var lazyConnectionField = typeof(RedisConnectionManager)
            .GetField("_lazyConnection", System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
        
        var lazyConnection = new Lazy<IConnectionMultiplexer>(() => conn);
        lazyConnectionField!.SetValue(manager, lazyConnection);
        
        return manager;
    }

    [Fact]
    public void GetConnection_ShouldReturnConnection()
    {
        // Arrange
        var manager = CreateManager();

        // Act
        var connection = manager.GetConnection();

        // Assert
        connection.Should().BeSameAs(_connectionMock.Object);
    }

    [Fact]
    public async Task AcquireLockAsync_WhenLockAvailable_ShouldAcquireLock()
    {
        // Arrange
        var db = _databaseMock;
        
        db.Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(true);

        var manager = CreateManager();

        // Act
        var lockObj = await manager.AcquireLockAsync("test-key", TimeSpan.FromSeconds(10));

        // Assert
        lockObj.Should().NotBeNull();
    }

    [Fact]
    public async Task AcquireLockAsync_WhenLockNotAvailable_ShouldRetry()
    {
        // Arrange
        var callCount = 0;
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(() =>
            {
                callCount++;
                return callCount > 3; // Succeed after 3 retries
            });

        var manager = CreateManager();

        // Act
        var lockObj = await manager.AcquireLockAsync("test-key", TimeSpan.FromSeconds(10));

        // Assert
        lockObj.Should().NotBeNull();
        callCount.Should().BeGreaterThan(1);
    }

    [Fact]
    public async Task AcquireLockAsync_WhenRetryExhausted_ShouldThrowTimeoutException()
    {
        // Arrange
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(false); // Always fail

        var manager = CreateManager();

        // Act & Assert
        await FluentActions.Invoking(() => manager.AcquireLockAsync("test-key", TimeSpan.FromMilliseconds(100)))
            .Should().ThrowAsync<TimeoutException>();
    }

    [Fact]
    public async Task ExecuteInLockAsync_ShouldExecuteFunction()
    {
        // Arrange
        var functionExecuted = false;
        
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(true);

        var manager = CreateManager();

        // Act
        var result = await manager.ExecuteInLockAsync("test-key", () =>
        {
            functionExecuted = true;
            return Task.FromResult("success");
        });

        // Assert
        functionExecuted.Should().BeTrue();
        result.Should().Be("success");
    }

    [Fact]
    public async Task ExecuteInLockAsync_ShouldReleaseLockAfterExecution()
    {
        // Arrange
        var lockReleased = false;
        
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(true);

        _databaseMock
            .Setup(x => x.ScriptEvaluate(
                It.IsAny<string>(),
                It.IsAny<RedisKey[]>(),
                It.IsAny<RedisValue[]>(),
                It.IsAny<CommandFlags>()))
            .Callback(() => lockReleased = true);

        var manager = CreateManager();

        // Act
        await manager.ExecuteInLockAsync("test-key", () => Task.FromResult("done"));

        // Assert
        lockReleased.Should().BeTrue();
    }

    [Fact]
    public async Task AcquireLockAsync_ShouldUseCorrectKeyFormat()
    {
        // Arrange
        RedisKey? capturedKey = null;
        
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .Callback<RedisKey, RedisValue, TimeSpan?, bool, When, CommandFlags>((key, _, _, _, _, _) => 
            {
                capturedKey = key;
            })
            .Returns(Task.FromResult(true));


        var manager = CreateManager();

        // Act
        await manager.AcquireLockAsync("my-resource");

        // Assert
        capturedKey.Should().NotBeNull();
        capturedKey!.ToString().Should().Contain("lock:test-instance:");
        capturedKey.ToString().Should().Contain("my-resource");
    }

    [Fact]
    public async Task AcquireLockAsync_ShouldUseDefaultExpiryWhenNotProvided()
    {
        // Arrange
        TimeSpan? capturedExpiry = null;
        
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .Callback<RedisKey, RedisValue, TimeSpan?, bool, When, CommandFlags>((_, _, expiry, _, _, _) => 
            {
                capturedExpiry = expiry;
            })
            .Returns(Task.FromResult(true));


        var manager = CreateManager();

        // Act
        await manager.AcquireLockAsync("test-key");

        // Assert
        capturedExpiry.Should().NotBeNull();
        capturedExpiry.Should().Be(TimeSpan.FromSeconds(10)); // Default is 10 seconds
    }

    [Fact]
    public async Task ExecuteInLockAsync_WithException_ShouldStillReleaseLock()
    {
        // Arrange
        var lockReleased = false;
        
        _databaseMock
            .Setup(x => x.StringSetAsync(
                It.IsAny<RedisKey>(),
                It.IsAny<RedisValue>(),
                It.IsAny<TimeSpan?>(),
                It.IsAny<bool>(),
                When.NotExists,
                It.IsAny<CommandFlags>()))
            .ReturnsAsync(true);

        _databaseMock
            .Setup(x => x.ScriptEvaluate(
                It.IsAny<string>(),
                It.IsAny<RedisKey[]>(),
                It.IsAny<RedisValue[]>(),
                It.IsAny<CommandFlags>()))
            .Callback(() => lockReleased = true);

        var manager = CreateManager();

        // Act & Assert
        await FluentActions.Invoking(() => 
            manager.ExecuteInLockAsync<string>("test-key", () => throw new InvalidOperationException("Test")))
            .Should().ThrowAsync<InvalidOperationException>();
        await FluentActions.Invoking(() => 
            manager.ExecuteInLockAsync<string>("test-key", () => throw new InvalidOperationException("Test")))
            .Should().ThrowAsync<InvalidOperationException>();
        // Lock should still be released
        lockReleased.Should().BeTrue();
    }
}