ChatBot/ChatBot.Tests/Models/ChatSessionTests.cs

using ChatBot.Models;
using ChatBot.Models.Dto;
using FluentAssertions;
using OllamaSharp.Models.Chat;

namespace ChatBot.Tests.Models;

public class ChatSessionTests
{
    [Fact]
    public void Constructor_ShouldInitializeProperties()
    {
        // Arrange & Act
        var session = new ChatSession();

        // Assert
        session.ChatId.Should().Be(0);
        session.SessionId.Should().NotBeNullOrEmpty();
        session.ChatType.Should().Be("private");
        session.ChatTitle.Should().Be(string.Empty);
        session.Model.Should().Be(string.Empty);
        session.CreatedAt.Should().BeCloseTo(DateTime.UtcNow, TimeSpan.FromSeconds(1));
        session.LastUpdatedAt.Should().BeCloseTo(DateTime.UtcNow, TimeSpan.FromSeconds(1));
        session.GetAllMessages().Should().NotBeNull();
        session.GetAllMessages().Should().BeEmpty();
    }

    [Fact]
    public void AddUserMessage_ShouldAddMessageToSession()
    {
        // Arrange
        var session = new ChatSession();
        var content = "Hello";
        var username = "testuser";

        // Act
        session.AddUserMessage(content, username);

        // Assert
        session.GetAllMessages().Should().HaveCount(1);
        var message = session.GetAllMessages()[0];
        message.Role.Should().Be(ChatRole.User);
        message.Content.Should().Be(content);
    }

    [Fact]
    public void AddAssistantMessage_ShouldAddMessageToSession()
    {
        // Arrange
        var session = new ChatSession();
        var content = "Hello! How can I help you?";

        // Act
        session.AddAssistantMessage(content);

        // Assert
        session.GetAllMessages().Should().HaveCount(1);
        var message = session.GetAllMessages()[0];
        message.Role.Should().Be(ChatRole.Assistant);
        message.Content.Should().Be(content);
    }

    [Fact]
    public void AddMessage_ShouldAddMessageToSession()
    {
        // Arrange
        var session = new ChatSession();
        var content = "You are a helpful assistant.";
        var message = new ChatMessage { Role = ChatRole.System, Content = content };

        // Act
        session.AddMessage(message);

        // Assert
        session.GetAllMessages().Should().HaveCount(1);
        var addedMessage = session.GetAllMessages()[0];
        addedMessage.Role.Should().Be(ChatRole.System);
        addedMessage.Content.Should().Be(content);
    }

    [Fact]
    public void ClearHistory_ShouldRemoveAllMessages()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Hello", "testuser");
        session.AddAssistantMessage("Hi there!");

        // Act
        session.ClearHistory();

        // Assert
        session.GetAllMessages().Should().BeEmpty();
    }

    [Fact]
    public void GetMessageCount_ShouldReturnCorrectCount()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Hello", "testuser");
        session.AddAssistantMessage("Hi there!");
        session.AddUserMessage("How are you?", "testuser");

        // Act
        var count = session.GetMessageCount();

        // Assert
        count.Should().Be(3);
    }

    [Fact]
    public void GetLastMessage_ShouldReturnLastMessage()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Hello", "testuser");
        session.AddAssistantMessage("Hi there!");
        session.AddUserMessage("How are you?", "testuser");

        // Act
        var messages = session.GetAllMessages();
        var lastMessage = messages.LastOrDefault();

        // Assert
        lastMessage.Should().NotBeNull();
        lastMessage!.Content.Should().Be("How are you?");
        lastMessage.Role.Should().Be(ChatRole.User);
    }

    [Fact]
    public void GetLastMessage_ShouldReturnNull_WhenNoMessages()
    {
        // Arrange
        var session = new ChatSession();

        // Act
        var messages = session.GetAllMessages();
        var lastMessage = messages.LastOrDefault();

        // Assert
        lastMessage.Should().BeNull();
    }

    [Fact]
    public void AddMessage_ShouldUpdateLastUpdatedAt()
    {
        // Arrange
        var session = new ChatSession();
        var originalTime = session.LastUpdatedAt;
        var message = new ChatMessage { Role = ChatRole.User, Content = "Test" };

        // Act
        session.AddMessage(message);

        // Assert
        session.LastUpdatedAt.Should().BeAfter(originalTime);
    }

    [Fact]
    public void GetMessageCount_ShouldReturnZero_WhenNoMessages()
    {
        // Arrange
        var session = new ChatSession();

        // Act
        var count = session.GetMessageCount();

        // Assert
        count.Should().Be(0);
    }

    [Fact]
    public void GetMessageCount_ShouldReturnCorrectCount_WhenHasMessages()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Hello", "testuser");

        // Act
        var count = session.GetMessageCount();

        // Assert
        count.Should().Be(1);
    }

    [Fact]
    public void AddUserMessage_InGroupChat_ShouldIncludeUsername()
    {
        // Arrange
        var session = new ChatSession { ChatType = "group" };
        var content = "Hello";
        var username = "testuser";

        // Act
        session.AddUserMessage(content, username);

        // Assert
        var message = session.GetAllMessages()[0];
        message.Content.Should().Be("testuser: Hello");
    }

    [Fact]
    public void AddUserMessage_InPrivateChat_ShouldNotIncludeUsername()
    {
        // Arrange
        var session = new ChatSession { ChatType = "private" };
        var content = "Hello";
        var username = "testuser";

        // Act
        session.AddUserMessage(content, username);

        // Assert
        var message = session.GetAllMessages()[0];
        message.Content.Should().Be("Hello");
    }

    [Fact]
    public void AddMessage_ShouldTrimHistory_WhenExceedsMaxHistoryLength()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 5 };

        // Add system message
        session.AddMessage(new ChatMessage { Role = ChatRole.System, Content = "System prompt" });

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

        // Act & Assert
        session.GetMessageCount().Should().BeLessThanOrEqualTo(5);
        // System message should be preserved
        session.GetAllMessages()[0].Role.Should().Be(ChatRole.System);
    }

    [Fact]
    public void AddMessage_ShouldTrimHistory_WithoutSystemMessage()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 3 };

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

        // Act & Assert
        session.GetMessageCount().Should().BeLessThanOrEqualTo(3);
    }

    [Fact]
    public void SetCompressionService_ShouldSetService()
    {
        // Arrange
        var session = new ChatSession();
        var compressionServiceMock =
            new Moq.Mock<ChatBot.Services.Interfaces.IHistoryCompressionService>();

        // Act
        session.SetCompressionService(compressionServiceMock.Object);

        // Assert
        // The service should be set (no exception)
        session.Should().NotBeNull();
    }

    [Fact]
    public async Task AddMessageWithCompressionAsync_WithoutCompressionService_ShouldUseTrimming()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 3 };

        // Act
        for (int i = 0; i < 5; i++)
        {
            await session.AddMessageWithCompressionAsync(
                new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" },
                10,
                5
            );
        }

        // Assert
        session.GetMessageCount().Should().BeLessThanOrEqualTo(3);
    }

    [Fact]
    public async Task AddUserMessageWithCompressionAsync_ShouldAddMessage()
    {
        // Arrange
        var session = new ChatSession { ChatType = "private" };

        // Act
        await session.AddUserMessageWithCompressionAsync("Test message", "user", 10, 5);

        // Assert
        session.GetMessageCount().Should().Be(1);
        session.GetAllMessages()[0].Content.Should().Be("Test message");
    }

    [Fact]
    public async Task AddUserMessageWithCompressionAsync_InGroupChat_ShouldIncludeUsername()
    {
        // Arrange
        var session = new ChatSession { ChatType = "group" };

        // Act
        await session.AddUserMessageWithCompressionAsync("Test message", "user", 10, 5);

        // Assert
        session.GetMessageCount().Should().Be(1);
        session.GetAllMessages()[0].Content.Should().Be("user: Test message");
    }

    [Fact]
    public async Task AddAssistantMessageWithCompressionAsync_ShouldAddMessage()
    {
        // Arrange
        var session = new ChatSession();

        // Act
        await session.AddAssistantMessageWithCompressionAsync("Test response", 10, 5);

        // Assert
        session.GetMessageCount().Should().Be(1);
        session.GetAllMessages()[0].Role.Should().Be(ChatRole.Assistant);
        session.GetAllMessages()[0].Content.Should().Be("Test response");
    }

    [Fact]
    public async Task AddMessageWithCompressionAsync_ShouldTriggerTrimming_WhenNoCompressionService()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 2 };

        // Act
        for (int i = 0; i < 4; i++)
        {
            await session.AddMessageWithCompressionAsync(
                new ChatMessage { Role = ChatRole.User, Content = $"Message {i}" },
                2,
                1
            );
        }

        // Assert
        session.GetMessageCount().Should().BeLessThanOrEqualTo(2);
    }

    [Fact]
    public async Task ClearHistory_ShouldUpdateLastUpdatedAt()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Test", "user");
        var lastUpdated = session.LastUpdatedAt;
        await Task.Delay(10); // Small delay

        // Act
        session.ClearHistory();

        // Assert
        session.LastUpdatedAt.Should().BeAfter(lastUpdated);
    }

    [Fact]
    public void SetCreatedAtForTesting_ShouldUpdateCreatedAt()
    {
        // Arrange
        var session = new ChatSession();
        var targetDate = DateTime.UtcNow.AddDays(-5);

        // Act
        session.SetCreatedAtForTesting(targetDate);

        // Assert
        session.CreatedAt.Should().Be(targetDate);
    }

    [Fact]
    public void AddMessage_MultipleTimes_ShouldMaintainOrder()
    {
        // Arrange
        var session = new ChatSession();

        // Act
        session.AddUserMessage("Message 1", "user1");
        session.AddAssistantMessage("Response 1");
        session.AddUserMessage("Message 2", "user1");
        session.AddAssistantMessage("Response 2");

        // Assert
        var messages = session.GetAllMessages();
        messages.Should().HaveCount(4);
        messages[0].Content.Should().Be("Message 1");
        messages[1].Content.Should().Be("Response 1");
        messages[2].Content.Should().Be("Message 2");
        messages[3].Content.Should().Be("Response 2");
    }

    [Fact]
    public void AddMessage_WithSystemMessage_ShouldPreserveSystemMessage()
    {
        // Arrange
        var session = new ChatSession { MaxHistoryLength = 3 };
        session.AddMessage(new ChatMessage { Role = ChatRole.System, Content = "System prompt" });

        // Act
        for (int i = 0; i < 5; i++)
        {
            session.AddUserMessage($"Message {i}", "user");
        }

        // Assert
        var messages = session.GetAllMessages();
        messages[0].Role.Should().Be(ChatRole.System);
        messages[0].Content.Should().Be("System prompt");
    }

    [Fact]
    public void GetAllMessages_ShouldReturnCopy()
    {
        // Arrange
        var session = new ChatSession();
        session.AddUserMessage("Test", "user");

        // Act
        var messages1 = session.GetAllMessages();
        var messages2 = session.GetAllMessages();

        // Assert
        messages1.Should().NotBeSameAs(messages2);
        messages1.Should().BeEquivalentTo(messages2);
    }

    [Fact]
    public async Task ChatSession_ThreadSafety_MultipleConcurrentAdds()
    {
        // Arrange
        var session = new ChatSession();
        var tasks = new List<Task>();

        // Act
        for (int i = 0; i < 10; i++)
        {
            int messageNum = i;
            tasks.Add(Task.Run(() => session.AddUserMessage($"Message {messageNum}", "user")));
        }

        await Task.WhenAll(tasks.ToArray());

        // Assert
        session.GetMessageCount().Should().Be(10);
    }

    [Fact]
    public void MaxHistoryLength_ShouldBeSettable()
    {
        // Arrange & Act
        var session = new ChatSession { MaxHistoryLength = 50 };

        // Assert
        session.MaxHistoryLength.Should().Be(50);
    }

    [Fact]
    public void Model_ShouldBeSettable()
    {
        // Arrange & Act
        var session = new ChatSession { Model = "llama3" };

        // Assert
        session.Model.Should().Be("llama3");
    }

    [Fact]
    public void SessionProperties_ShouldBeSettableViaInitializer()
    {
        // Arrange & Act
        var session = new ChatSession
        {
            ChatId = 12345,
            ChatType = "group",
            ChatTitle = "Test Group",
            Model = "llama3",
            MaxHistoryLength = 50,
        };

        // Assert
        session.ChatId.Should().Be(12345);
        session.ChatType.Should().Be("group");
        session.ChatTitle.Should().Be("Test Group");
        session.Model.Should().Be("llama3");
        session.MaxHistoryLength.Should().Be(50);
    }
}