ChatBot/ChatBot.Tests/Services/DatabaseSessionStorageTests.cs

using ChatBot.Data;
using ChatBot.Data.Interfaces;
using ChatBot.Models.Entities;
using ChatBot.Services;
using ChatBot.Tests.TestUtilities;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Moq;

namespace ChatBot.Tests.Services;

public class DatabaseSessionStorageTests : TestBase
{
    private ChatBotDbContext _dbContext = null!;
    private DatabaseSessionStorage _sessionStorage = null!;
    private Mock<IChatSessionRepository> _repositoryMock = null!;

    public DatabaseSessionStorageTests()
    {
        SetupServices();
    }

    protected override void ConfigureServices(IServiceCollection services)
    {
        // Add in-memory database
        services.AddDbContext<ChatBotDbContext>(options =>
            options.UseInMemoryDatabase("TestDatabase")
        );

        // Add mocked repository
        _repositoryMock = new Mock<IChatSessionRepository>();
        services.AddSingleton(_repositoryMock.Object);

        // Add logger
        services.AddSingleton(Mock.Of<ILogger<DatabaseSessionStorage>>());

        // Add session storage
        services.AddScoped<DatabaseSessionStorage>();
    }

    protected override void SetupServices()
    {
        base.SetupServices();

        _dbContext = ServiceProvider.GetRequiredService<ChatBotDbContext>();
        _sessionStorage = ServiceProvider.GetRequiredService<DatabaseSessionStorage>();

        // Ensure database is created
        _dbContext.Database.EnsureCreated();
    }

    [Fact]
    public void GetOrCreate_ShouldReturnExistingSession_WhenSessionExists()
    {
        // Arrange
        var existingSession = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock
            .Setup(x => x.GetOrCreateAsync(12345, "private", "Test Chat"))
            .ReturnsAsync(existingSession);

        // Act
        var result = _sessionStorage.GetOrCreate(12345, "private", "Test Chat");

        // Assert
        result.Should().NotBeNull();
        result.ChatId.Should().Be(12345);
        _repositoryMock.Verify(x => x.GetOrCreateAsync(12345, "private", "Test Chat"), Times.Once);
    }

    [Fact]
    public void Get_ShouldReturnSession_WhenSessionExists()
    {
        // Arrange
        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock.Setup(x => x.GetByChatIdAsync(12345)).ReturnsAsync(sessionEntity);

        // Act
        var result = _sessionStorage.Get(12345);

        // Assert
        result.Should().NotBeNull();
        result.ChatId.Should().Be(12345);
        _repositoryMock.Verify(x => x.GetByChatIdAsync(12345), Times.Once);
    }

    [Fact]
    public void Get_ShouldReturnNull_WhenSessionDoesNotExist()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.GetByChatIdAsync(12345))
            .ReturnsAsync((ChatSessionEntity?)null);

        // Act
        var result = _sessionStorage.Get(12345);

        // Assert
        result.Should().BeNull();
        _repositoryMock.Verify(x => x.GetByChatIdAsync(12345), Times.Once);
    }

    [Fact]
    public async Task SaveSessionAsync_ShouldUpdateSession_WhenSessionExists()
    {
        // Arrange
        var session = TestDataBuilder.ChatSessions.CreateBasicSession(12345, "private");
        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock.Setup(x => x.GetByChatIdAsync(12345)).ReturnsAsync(sessionEntity);
        _repositoryMock
            .Setup(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()))
            .ReturnsAsync(sessionEntity);

        // Act
        await _sessionStorage.SaveSessionAsync(session);

        // Assert
        _repositoryMock.Verify(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()), Times.Once);
    }

    [Fact]
    public void Remove_ShouldReturnTrue_WhenSessionExists()
    {
        // Arrange
        _repositoryMock.Setup(x => x.DeleteAsync(12345)).ReturnsAsync(true);

        // Act
        var result = _sessionStorage.Remove(12345);

        // Assert
        result.Should().BeTrue();
        _repositoryMock.Verify(x => x.DeleteAsync(12345), Times.Once);
    }

    [Fact]
    public void Remove_ShouldReturnFalse_WhenSessionDoesNotExist()
    {
        // Arrange
        _repositoryMock.Setup(x => x.DeleteAsync(12345)).ReturnsAsync(false);

        // Act
        var result = _sessionStorage.Remove(12345);

        // Assert
        result.Should().BeFalse();
        _repositoryMock.Verify(x => x.DeleteAsync(12345), Times.Once);
    }

    [Fact]
    public void GetActiveSessionsCount_ShouldReturnCorrectCount()
    {
        // Arrange
        var expectedCount = 5;
        _repositoryMock.Setup(x => x.GetActiveSessionsCountAsync()).ReturnsAsync(expectedCount);

        // Act
        var result = _sessionStorage.GetActiveSessionsCount();

        // Assert
        result.Should().Be(expectedCount);
        _repositoryMock.Verify(x => x.GetActiveSessionsCountAsync(), Times.Once);
    }

    [Fact]
    public void CleanupOldSessions_ShouldReturnCorrectCount()
    {
        // Arrange
        var expectedCount = 3;
        _repositoryMock.Setup(x => x.CleanupOldSessionsAsync(24)).ReturnsAsync(expectedCount);

        // Act
        var result = _sessionStorage.CleanupOldSessions(24);

        // Assert
        result.Should().Be(expectedCount);
        _repositoryMock.Verify(x => x.CleanupOldSessionsAsync(24), Times.Once);
    }

    [Fact]
    public void GetOrCreate_ShouldThrowInvalidOperationException_WhenRepositoryThrows()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.GetOrCreateAsync(12345, "private", "Test Chat"))
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var act = () => _sessionStorage.GetOrCreate(12345, "private", "Test Chat");

        // Assert
        act.Should()
            .Throw<InvalidOperationException>()
            .WithMessage("Failed to get or create session for chat 12345")
            .WithInnerException<Exception>()
            .WithMessage("Database error");
    }

    [Fact]
    public void Get_ShouldReturnNull_WhenRepositoryThrows()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.GetByChatIdAsync(12345))
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var result = _sessionStorage.Get(12345);

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

    [Fact]
    public async Task SaveSessionAsync_ShouldLogWarning_WhenSessionNotFound()
    {
        // Arrange
        var session = TestDataBuilder.ChatSessions.CreateBasicSession(12345, "private");
        _repositoryMock
            .Setup(x => x.GetByChatIdAsync(12345))
            .ReturnsAsync((ChatSessionEntity?)null);

        // Act
        await _sessionStorage.SaveSessionAsync(session);

        // Assert
        _repositoryMock.Verify(x => x.GetByChatIdAsync(12345), Times.Once);
        _repositoryMock.Verify(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()), Times.Never);
    }

    [Fact]
    public async Task SaveSessionAsync_ShouldThrowInvalidOperationException_WhenRepositoryThrows()
    {
        // Arrange
        var session = TestDataBuilder.ChatSessions.CreateBasicSession(12345, "private");
        _repositoryMock
            .Setup(x => x.GetByChatIdAsync(12345))
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var act = async () => await _sessionStorage.SaveSessionAsync(session);

        // Assert
        var exception = await act.Should()
            .ThrowAsync<InvalidOperationException>()
            .WithMessage("Failed to save session for chat 12345");
        exception
            .And.InnerException.Should()
            .BeOfType<Exception>()
            .Which.Message.Should()
            .Be("Database error");
    }

    [Fact]
    public async Task SaveSessionAsync_ShouldClearMessagesAndAddNew()
    {
        // Arrange
        var session = TestDataBuilder.ChatSessions.CreateBasicSession(12345, "private");
        session.AddUserMessage("Test message", "user1");
        session.AddAssistantMessage("Test response");

        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock.Setup(x => x.GetByChatIdAsync(12345)).ReturnsAsync(sessionEntity);
        _repositoryMock
            .Setup(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()))
            .ReturnsAsync(sessionEntity);

        // Act
        await _sessionStorage.SaveSessionAsync(session);

        // Assert
        _repositoryMock.Verify(x => x.ClearMessagesAsync(It.IsAny<int>()), Times.Once);
        _repositoryMock.Verify(
            x =>
                x.AddMessageAsync(
                    It.IsAny<int>(),
                    It.IsAny<string>(),
                    It.IsAny<string>(),
                    It.IsAny<int>()
                ),
            Times.Exactly(2)
        );
        _repositoryMock.Verify(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()), Times.Once);
    }

    [Fact]
    public void Remove_ShouldReturnFalse_WhenRepositoryThrows()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.DeleteAsync(12345))
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var result = _sessionStorage.Remove(12345);

        // Assert
        result.Should().BeFalse();
    }

    [Fact]
    public void GetActiveSessionsCount_ShouldReturnZero_WhenRepositoryThrows()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.GetActiveSessionsCountAsync())
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var result = _sessionStorage.GetActiveSessionsCount();

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

    [Fact]
    public void CleanupOldSessions_ShouldReturnZero_WhenRepositoryThrows()
    {
        // Arrange
        _repositoryMock
            .Setup(x => x.CleanupOldSessionsAsync(24))
            .ThrowsAsync(new Exception("Database error"));

        // Act
        var result = _sessionStorage.CleanupOldSessions(24);

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

    [Fact]
    public void GetOrCreate_WithCompressionService_ShouldSetCompressionService()
    {
        // Arrange
        var compressionServiceMock = TestDataBuilder.Mocks.CreateCompressionServiceMock();
        var storageWithCompression = new DatabaseSessionStorage(
            _repositoryMock.Object,
            Mock.Of<ILogger<DatabaseSessionStorage>>(),
            compressionServiceMock.Object
        );

        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock
            .Setup(x => x.GetOrCreateAsync(12345, "private", "Test Chat"))
            .ReturnsAsync(sessionEntity);

        // Act
        var result = storageWithCompression.GetOrCreate(12345, "private", "Test Chat");

        // Assert
        result.Should().NotBeNull();
        result.ChatId.Should().Be(12345);
    }

    [Fact]
    public void Get_WithCompressionService_ShouldSetCompressionService()
    {
        // Arrange
        var loggerMock = new Mock<ILogger<DatabaseSessionStorage>>();
        var compressionServiceMock = TestDataBuilder.Mocks.CreateCompressionServiceMock();
        var storageWithCompression = new DatabaseSessionStorage(
            _repositoryMock.Object,
            loggerMock.Object,
            compressionServiceMock.Object
        );

        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        sessionEntity.Messages.Add(
            new ChatMessageEntity
            {
                Id = 1,
                SessionId = sessionEntity.Id,
                Content = "Test",
                Role = "user",
                MessageOrder = 0,
                CreatedAt = DateTime.UtcNow,
            }
        );

        _repositoryMock.Setup(x => x.GetByChatIdAsync(12345)).ReturnsAsync(sessionEntity);

        // Act
        var result = storageWithCompression.Get(12345);

        // Assert
        _repositoryMock.Verify(x => x.GetByChatIdAsync(12345), Times.Once);
        result.Should().NotBeNull();
        result!.GetMessageCount().Should().Be(1);
    }

    [Fact]
    public async Task SaveSessionAsync_WithMultipleMessages_ShouldSaveInCorrectOrder()
    {
        // Arrange
        var session = TestDataBuilder.ChatSessions.CreateBasicSession(12345, "private");
        session.AddUserMessage("Message 1", "user1");
        session.AddAssistantMessage("Response 1");
        session.AddUserMessage("Message 2", "user1");
        session.AddAssistantMessage("Response 2");

        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock.Setup(x => x.GetByChatIdAsync(12345)).ReturnsAsync(sessionEntity);
        _repositoryMock
            .Setup(x => x.UpdateAsync(It.IsAny<ChatSessionEntity>()))
            .ReturnsAsync(sessionEntity);

        // Act
        await _sessionStorage.SaveSessionAsync(session);

        // Assert
        _repositoryMock.Verify(x => x.ClearMessagesAsync(It.IsAny<int>()), Times.Once);
        _repositoryMock.Verify(
            x =>
                x.AddMessageAsync(
                    It.IsAny<int>(),
                    It.IsAny<string>(),
                    It.IsAny<string>(),
                    It.IsAny<int>()
                ),
            Times.Exactly(4)
        );
    }

    [Fact]
    public void GetOrCreate_WithDefaultParameters_ShouldUseDefaults()
    {
        // Arrange
        var sessionEntity = TestDataBuilder.Mocks.CreateChatSessionEntity();
        _repositoryMock
            .Setup(x => x.GetOrCreateAsync(12345, "private", ""))
            .ReturnsAsync(sessionEntity);

        // Act
        var result = _sessionStorage.GetOrCreate(12345);

        // Assert
        result.Should().NotBeNull();
        _repositoryMock.Verify(x => x.GetOrCreateAsync(12345, "private", ""), Times.Once);
    }
}