using ChatBot.Models; using ChatBot.Services.Interfaces; using FluentAssertions; using Moq; using OllamaSharp.Models.Chat; namespace ChatBot.Tests.Models; public class ChatSessionCompressionTests { [Fact] public async Task CompressHistoryAsync_ShouldCompressMessages_WhenCompressionServiceAvailable() { // Arrange var session = new ChatSession(); var compressionServiceMock = new Mock(); session.SetCompressionService(compressionServiceMock.Object); // Setup compression service to return compressed messages compressionServiceMock .Setup(x => x.CompressHistoryAsync(It.IsAny>(), It.IsAny())) .ReturnsAsync( new List { new() { Role = ChatRole.System, Content = "System prompt" }, new() { Role = ChatRole.User, Content = "Compressed user message" }, } ); // 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(); session.SetCompressionService(compressionServiceMock.Object); // Setup compression service to throw an exception compressionServiceMock .Setup(x => x.CompressHistoryAsync(It.IsAny>(), It.IsAny())) .ThrowsAsync(new Exception("Compression failed")); // 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(); session.SetCompressionService(compressionServiceMock.Object); // 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>(), It.IsAny()), 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(); session.SetCompressionService(compressionServiceMock.Object); // Setup compression service to delay to simulate processing time compressionServiceMock .Setup(x => x.CompressHistoryAsync(It.IsAny>(), It.IsAny())) .ReturnsAsync( (List messages, int target) => Task.Delay(50) .ContinueWith(_ => new List { new() { Role = ChatRole.System, Content = "Compressed" }, }) ); var tasks = new List(); 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().Object; var compressionService2 = new Mock().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(); session.SetCompressionService(compressionServiceMock.Object); // Setup compression service to preserve system message compressionServiceMock .Setup(x => x.CompressHistoryAsync(It.IsAny>(), It.IsAny())) .ReturnsAsync( (List messages, int target) => { var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System); var compressed = new List(); if (systemMessage != null) { compressed.Add(systemMessage); } compressed.Add( new ChatMessage { Role = ChatRole.User, Content = "Compressed user messages", } ); return compressed; } ); // 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"); } }