ColaFlow/colaflow-api/tests/Modules/Mcp/ColaFlow.Modules.Mcp.Tests/Handlers/McpProtocolHandlerTests.cs

using ColaFlow.Modules.Mcp.Application.Handlers;
using ColaFlow.Modules.Mcp.Contracts.JsonRpc;
using FluentAssertions;
using Microsoft.Extensions.Logging;
using NSubstitute;

namespace ColaFlow.Modules.Mcp.Tests.Handlers;

/// <summary>
/// Unit tests for McpProtocolHandler
/// </summary>
public class McpProtocolHandlerTests
{
    private readonly ILogger<McpProtocolHandler> _logger;
    private readonly List<IMcpMethodHandler> _methodHandlers;
    private readonly McpProtocolHandler _sut;

    public McpProtocolHandlerTests()
    {
        _logger = Substitute.For<ILogger<McpProtocolHandler>>();
        _methodHandlers = new List<IMcpMethodHandler>();
        _sut = new McpProtocolHandler(_logger, _methodHandlers);
    }

    [Fact]
    public async Task HandleRequestAsync_WithInvalidJsonRpc_ReturnsInvalidRequest()
    {
        // Arrange
        var request = new JsonRpcRequest
        {
            JsonRpc = "1.0", // Invalid version
            Method = "test"
        };

        // Act
        var response = await _sut.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.JsonRpc.Should().Be("2.0");
        response.Error.Should().NotBeNull();
        response.Error!.Code.Should().Be((int)JsonRpcErrorCode.InvalidRequest);
        response.Result.Should().BeNull();
    }

    [Fact]
    public async Task HandleRequestAsync_WithMissingMethod_ReturnsInvalidRequest()
    {
        // Arrange
        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "" // Empty method
        };

        // Act
        var response = await _sut.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.Error.Should().NotBeNull();
        response.Error!.Code.Should().Be((int)JsonRpcErrorCode.InvalidRequest);
    }

    [Fact]
    public async Task HandleRequestAsync_WithUnknownMethod_ReturnsMethodNotFound()
    {
        // Arrange
        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "unknown_method",
            Id = 1
        };

        // Act
        var response = await _sut.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.JsonRpc.Should().Be("2.0");
        response.Id.Should().Be(1);
        response.Error.Should().NotBeNull();
        response.Error!.Code.Should().Be((int)JsonRpcErrorCode.MethodNotFound);
        response.Error.Message.Should().Contain("unknown_method");
        response.Result.Should().BeNull();
    }

    [Fact]
    public async Task HandleRequestAsync_WithValidMethod_ReturnsSuccess()
    {
        // Arrange
        var mockHandler = Substitute.For<IMcpMethodHandler>();
        mockHandler.MethodName.Returns("test_method");
        mockHandler.HandleAsync(Arg.Any<object?>(), Arg.Any<CancellationToken>())
            .Returns(Task.FromResult<object?>(new { result = "success" }));

        var handler = new McpProtocolHandler(_logger, new[] { mockHandler });

        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "test_method",
            Id = 1
        };

        // Act
        var response = await handler.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.JsonRpc.Should().Be("2.0");
        response.Id.Should().Be(1);
        response.Error.Should().BeNull();
        response.Result.Should().NotBeNull();
    }

    [Fact]
    public async Task HandleRequestAsync_WhenHandlerThrowsArgumentException_ReturnsInvalidParams()
    {
        // Arrange
        var mockHandler = Substitute.For<IMcpMethodHandler>();
        mockHandler.MethodName.Returns("test_method");
        mockHandler.HandleAsync(Arg.Any<object?>(), Arg.Any<CancellationToken>())
            .Returns<object?>(_ => throw new ArgumentException("Invalid parameter"));

        var handler = new McpProtocolHandler(_logger, new[] { mockHandler });

        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "test_method",
            Id = 1
        };

        // Act
        var response = await handler.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.Error.Should().NotBeNull();
        response.Error!.Code.Should().Be((int)JsonRpcErrorCode.InvalidParams);
    }

    [Fact]
    public async Task HandleRequestAsync_WhenHandlerThrowsException_ReturnsInternalError()
    {
        // Arrange
        var mockHandler = Substitute.For<IMcpMethodHandler>();
        mockHandler.MethodName.Returns("test_method");
        mockHandler.HandleAsync(Arg.Any<object?>(), Arg.Any<CancellationToken>())
            .Returns<object?>(_ => throw new InvalidOperationException("Something went wrong"));

        var handler = new McpProtocolHandler(_logger, new[] { mockHandler });

        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "test_method",
            Id = 1
        };

        // Act
        var response = await handler.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Should().NotBeNull();
        response.Error.Should().NotBeNull();
        response.Error!.Code.Should().Be((int)JsonRpcErrorCode.InternalError);
    }

    [Fact]
    public async Task HandleRequestAsync_WithStringId_PreservesIdInResponse()
    {
        // Arrange
        var mockHandler = Substitute.For<IMcpMethodHandler>();
        mockHandler.MethodName.Returns("test_method");
        mockHandler.HandleAsync(Arg.Any<object?>(), Arg.Any<CancellationToken>())
            .Returns(Task.FromResult<object?>(new { result = "success" }));

        var handler = new McpProtocolHandler(_logger, new[] { mockHandler });

        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "test_method",
            Id = "abc-123"
        };

        // Act
        var response = await handler.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        response.Id.Should().Be("abc-123");
    }

    [Fact]
    public async Task HandleRequestAsync_WithParams_PassesParamsToHandler()
    {
        // Arrange
        var mockHandler = Substitute.For<IMcpMethodHandler>();
        mockHandler.MethodName.Returns("test_method");
        mockHandler.HandleAsync(Arg.Any<object?>(), Arg.Any<CancellationToken>())
            .Returns(Task.FromResult<object?>(new { result = "success" }));

        var handler = new McpProtocolHandler(_logger, new[] { mockHandler });

        var testParams = new { param1 = "value1", param2 = 42 };
        var request = new JsonRpcRequest
        {
            JsonRpc = "2.0",
            Method = "test_method",
            Params = testParams,
            Id = 1
        };

        // Act
        await handler.HandleRequestAsync(request, CancellationToken.None);

        // Assert
        await mockHandler.Received(1).HandleAsync(testParams, Arg.Any<CancellationToken>());
    }
}