ChatBot.Tests/Models/ChatSessionCompressionTests.cs

@@ -1,8 +1,10 @@
 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;
 
@@ -17,15 +19,17 @@ public class ChatSessionCompressionTests
         session.SetCompressionService(compressionServiceMock.Object);
 
         // Setup compression service to return compressed messages
-        compressionServiceMock
+        var compressedMessages = new List<ChatMessage>
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ReturnsAsync(
-                new List<ChatMessage>
         {
-                    new() { Role = ChatRole.System, Content = "System prompt" },
+            new ChatMessage { Role = ChatRole.System.ToString(), Content = "System prompt" },
-                    new() { Role = ChatRole.User, Content = "Compressed user message" },
+            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++)
@@ -57,9 +61,13 @@ public class ChatSessionCompressionTests
         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>()))
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
-            .ThrowsAsync(new Exception("Compression failed"));
+            .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++)
@@ -87,6 +95,11 @@ public class ChatSessionCompressionTests
         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" });
@@ -100,7 +113,7 @@ public class ChatSessionCompressionTests
 
         // Assert - Should not call compression service
         compressionServiceMock.Verify(
-            x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()),
+            x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()),
             Times.Never
         );
 
@@ -116,17 +129,21 @@ public class ChatSessionCompressionTests
         var compressionServiceMock = new Mock<IHistoryCompressionService>();
         session.SetCompressionService(compressionServiceMock.Object);
 
-        // Setup compression service to delay to simulate processing time
+        // Setup compression service to simulate processing time
-        compressionServiceMock
+        var delayedResult = new List<ChatMessage>
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
-            .ReturnsAsync(
-                (List<ChatMessage> messages, int target) =>
-                    Task.Delay(50)
-                        .ContinueWith(_ => new List<ChatMessage>
         {
-                            new() { Role = ChatRole.System, Content = "Compressed" },
+            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);
+                return delayedResult;
+            });
+        compressionServiceMock
+            .Setup(x => x.ShouldCompress(It.IsAny<int>(), It.IsAny<int>()))
+            .Returns(true);
 
         var tasks = new List<Task>();
         int messageCount = 5;
@@ -176,11 +193,10 @@ public class ChatSessionCompressionTests
 
         // Setup compression service to preserve system message
         compressionServiceMock
-            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
+            .Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(
+            .Returns((List<ChatMessage> messages, int target, CancellationToken ct) =>
-                (List<ChatMessage> messages, int target) =>
             {
-                    var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System);
+                var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System.ToString());
                 var compressed = new List<ChatMessage>();
 
                 if (systemMessage != null)
@@ -188,17 +204,17 @@ public class ChatSessionCompressionTests
                     compressed.Add(systemMessage);
                 }
 
-                    compressed.Add(
+                compressed.Add(new ChatMessage
-                        new ChatMessage
                 {
-                            Role = ChatRole.User,
+                    Role = ChatRole.User.ToString(),
-                            Content = "Compressed user messages",
+                    Content = "Compressed user messages"
-                        }
+                });
-                    );
 
-                    return compressed;
+                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" });