Add tests
This commit is contained in:
Leonid Pershin
224
ChatBot.Tests/Models/ChatSessionCompressionTests.cs
Normal file
224
ChatBot.Tests/Models/ChatSessionCompressionTests.cs
Normal file
@@ -0,0 +1,224 @@
|
|||||||
|
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<IHistoryCompressionService>();
|
||||||
|
session.SetCompressionService(compressionServiceMock.Object);
|
||||||
|
|
||||||
|
// Setup compression service to return compressed messages
|
||||||
|
compressionServiceMock
|
||||||
|
.Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
|
||||||
|
.ReturnsAsync(
|
||||||
|
new List<ChatMessage>
|
||||||
|
{
|
||||||
|
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<IHistoryCompressionService>();
|
||||||
|
session.SetCompressionService(compressionServiceMock.Object);
|
||||||
|
|
||||||
|
// Setup compression service to throw an exception
|
||||||
|
compressionServiceMock
|
||||||
|
.Setup(x => x.CompressHistoryAsync(It.IsAny<List<ChatMessage>>(), It.IsAny<int>()))
|
||||||
|
.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<IHistoryCompressionService>();
|
||||||
|
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<List<ChatMessage>>(), It.IsAny<int>()),
|
||||||
|
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 delay to simulate processing time
|
||||||
|
compressionServiceMock
|
||||||
|
.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" },
|
||||||
|
})
|
||||||
|
);
|
||||||
|
|
||||||
|
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>()))
|
||||||
|
.ReturnsAsync(
|
||||||
|
(List<ChatMessage> messages, int target) =>
|
||||||
|
{
|
||||||
|
var systemMessage = messages.FirstOrDefault(m => m.Role == ChatRole.System);
|
||||||
|
var compressed = new List<ChatMessage>();
|
||||||
|
|
||||||
|
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");
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user