ChatBot/ChatBot.Tests/Models/ChatSessionCompressionTests.cs

using ChatBot.Models;
using ChatBot.Models.Dto;
using ChatBot.Services.Interfaces;
using FluentAssertions;
using Moq;
using OllamaSharp.Models.Chat;
using System.Collections.Concurrent;

namespace ChatBot.Tests.Models;

public class ChatSessionCompressionTests
{
    [Fact]
    public async Task CompressHistoryAsync_ShouldCompressMessages_WhenCompressionServiceAvailable()
    {
        // Arrange
        var session = new ChatSession();
        var compressionServiceMock = new Mock<IHistoryCompressionService>();
        session.SetCompressionService(compressionServiceMock.Object);

        // Setup compression service to return compressed messages
        var compressedMessages = new List<ChatMessage>
        {
            new ChatMessage { Role = ChatRole.System.ToString(), Content = "System prompt" },
            new ChatMessage { Role = ChatRole.User.ToString(), Content = "Compressed user message" }
        };
        compressionServiceMock
            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
            .ReturnsAsync(compressedMessages);
        compressionServiceMock
            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
            .Returns(true);

        // Add messages to session
        for (int i = 0; i < 10; i++)
        {
            session.AddMessage(new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" });
        }

        // Act
        await session.AddMessageWithCompressionAsync(
            new ChatMessage { Role = ChatRole.User, Content = "New message" },
            compressionThreshold: 5,
            compressionTarget: 2
        );

        // Assert
        var messages = session.GetAllMessages();
        messages.Should().HaveCount(2);
        messages[0].Role.Should().Be(ChatRole.System);
        messages[1].Role.Should().Be(ChatRole.User);
        messages[1].Content.Should().Be("Compressed user message");
    }

    [Fact]
    public async Task CompressHistoryAsync_ShouldFallbackToTrimming_WhenCompressionFails()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 3 };
        var compressionServiceMock = new Mock<IHistoryCompressionService>();
        session.SetCompressionService(compressionServiceMock.Object);

        // Setup compression service to throw an exception
        var exception = new Exception("Compression failed");
        compressionServiceMock
            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
            .ThrowsAsync(exception);
        compressionServiceMock
            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
            .Returns(true);

        // Add messages to session
        for (int i = 0; i < 5; i++)
        {
            session.AddMessage(new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" });
        }

        // Act
        await session.AddMessageWithCompressionAsync(
            new ChatMessage { Role = ChatRole.User, Content = "New message" },
            compressionThreshold: 3,
            compressionTarget: 2
        );

        // Assert - Should fall back to simple trimming
        var messages = session.GetAllMessages();
        messages.Should().HaveCount(3);
    }

    [Fact]
    public async Task AddMessageWithCompressionAsync_ShouldNotCompress_WhenBelowThreshold()
    {
        // Arrange
        var session = new ChatSession();
        var compressionServiceMock = new Mock<IHistoryCompressionService>();
        session.SetCompressionService(compressionServiceMock.Object);
        
        // Setup compression service to return false for ShouldCompress when count is below threshold
        compressionServiceMock
            .Setup(x => x.ShouldCompress(It.Is<int>(c => c < 5), It.Is<int>(t => t == 5)))
            .Returns(false);

        // Add messages to session (below threshold)
        session.AddMessage(new ChatMessage { Role = ChatRole.User, Content = "Message 1" });
        session.AddMessage(new ChatMessage { Role = ChatRole.Assistant, Content = "Response 1" });

        // Act - Set threshold higher than current message count
        await session.AddMessageWithCompressionAsync(
            new ChatMessage { Role = ChatRole.User, Content = "Message 2" },
            compressionThreshold: 5,
            compressionTarget: 2
        );

        // Assert - Should not call compression service
        compressionServiceMock.Verify(
            x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()),
            Times.Never
        );

        var messages = session.GetAllMessages();
        messages.Should().HaveCount(3);
    }

    [Fact]
    public async Task AddMessageWithCompressionAsync_ShouldHandleConcurrentAccess()
    {
        // Arrange
        var session = new ChatSession();
        var compressionServiceMock = new Mock<IHistoryCompressionService>();
        session.SetCompressionService(compressionServiceMock.Object);

        // Setup compression service to simulate processing time
        var delayedResult = new List<ChatMessage>
        {
            new ChatMessage { Role = ChatRole.System.ToString(), Content = "Compressed" }
        };
        compressionServiceMock
            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
            .Returns(async (List<ChatMessage> messages, int target, CancellationToken ct) =>
            {
                await Task.Delay(50, ct);
                return delayedResult;
            });
        compressionServiceMock
            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
            .Returns(true);

        var tasks = new List<Task>();
        int messageCount = 5;

        // Act - Start multiple concurrent operations
        for (int i = 0; i < messageCount; i++)
        {
            tasks.Add(
                session.AddMessageWithCompressionAsync(
                    new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" },
                    compressionThreshold: 2,
                    compressionTarget: 1
                )
            );
        }

        // Wait for all operations to complete
        await Task.WhenAll(tasks);

        // Assert - Should handle concurrent access without exceptions
        // and maintain thread safety
        session.GetMessageCount().Should().Be(1);
    }

    [Fact]
    public void SetCompressionService_ShouldNotThrow_WhenCalledMultipleTimes()
    {
        // Arrange
        var session = new ChatSession();
        var compressionService1 = new Mock<IHistoryCompressionService>().Object;
        var compressionService2 = new Mock<IHistoryCompressionService>().Object;

        // Act & Assert
        session.Invoking(s => s.SetCompressionService(compressionService1)).Should().NotThrow();

        // Should not throw when setting a different service
        session.Invoking(s => s.SetCompressionService(compressionService2)).Should().NotThrow();
    }

    [Fact]
    public async Task CompressHistoryAsync_ShouldPreserveSystemMessage_WhenCompressing()
    {
        // Arrange
        var session = new ChatSession();
        var compressionServiceMock = new Mock<IHistoryCompressionService>();
        session.SetCompressionService(compressionServiceMock.Object);

        // Setup compression service to preserve system message
        compressionServiceMock
            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
            .Returns((List<ChatMessage> messages, int target, CancellationToken ct) =>
            {
                var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System.ToString());
                var compressed = new List<ChatMessage>();

                if (systemMessage != null)
                {
                    compressed.Add(systemMessage);
                }

                compressed.Add(new ChatMessage
                {
                    Role = ChatRole.User.ToString(),
                    Content = "Compressed user messages"
                });

                return Task.FromResult(compressed);
            });
        compressionServiceMock
            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
            .Returns(true);

        // Add system message and some user messages
        session.AddMessage(new ChatMessage { Role = ChatRole.System, Content = "System prompt" });
        for (int i = 0; i < 10; i++)
        {
            session.AddMessage(new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" });
        }

        // Act
        await session.AddMessageWithCompressionAsync(
            new ChatMessage { Role = ChatRole.User, Content = "New message" },
            compressionThreshold: 5,
            compressionTarget: 2
        );

        // Assert - System message should be preserved
        var messages = session.GetAllMessages();
        messages.Should().HaveCount(2);
        messages[0].Role.Should().Be(ChatRole.System);
        messages[0].Content.Should().Be("System prompt");
        messages[1].Content.Should().Be("Compressed user messages");
    }
}