free-code-dotnet/tests/FreeCode.Tests.Unit/Engine/SystemPromptBuilderTests.cs

using Xunit;
using FluentAssertions;
using NSubstitute;
using System.Text.Json;
using FreeCode.Core.Enums;
using FreeCode.Core.Interfaces;
using FreeCode.Core.Models;
using FreeCode.Engine;
using FreeCode.Tests.Unit.Helpers;

namespace FreeCode.Tests.Unit.Engine;

public sealed class SystemPromptBuilderTests
{
    [Fact]
    public async Task BuildAsync_IncludesToolsCommandsAndMemory()
    {
        var sut = CreateBuilder(
            toolRegistry: new StubToolRegistry
            {
                Tools =
                [
                    new StubTool
                    {
                        Name = "Read",
                        Category = ToolCategory.FileSystem,
                        DescriptionFactory = _ => Task.FromResult("Reads files"),
                        InputSchema = ParseJson("{\"type\":\"object\"}")
                    }
                ]
            },
            commandRegistry: new StubCommandRegistry
            {
                EnabledCommandsOverride = [new StubCommand("help", "Shows help")]
            },
            memoryService: new StubSessionMemoryService { CurrentMemory = "remember this" });

        var result = await sut.BuildAsync([CreateMessage("hi")], new ToolPermissionContext(), new SubmitMessageOptions(Model: "claude", QuerySource: "test"));

        result.Should().Contain("<tool_descriptions>");
        result.Should().Contain("- name: Read");
        result.Should().Contain("/help - Shows help");
        result.Should().Contain("remember this");
        result.Should().Contain("selected_model: claude");
        result.Should().Contain("query_source: test");
    }

    [Fact]
    public async Task BuildAsync_SortsToolsAlphabetically()
    {
        var sut = CreateBuilder(
            toolRegistry: new StubToolRegistry
            {
                Tools =
                [
                    new StubTool { Name = "Zeta", Category = ToolCategory.FileSystem, DescriptionFactory = _ => Task.FromResult("z"), InputSchema = ParseJson("{\"type\":\"object\"}") },
                    new StubTool { Name = "Alpha", Category = ToolCategory.FileSystem, DescriptionFactory = _ => Task.FromResult("a"), InputSchema = ParseJson("{\"type\":\"object\"}") }
                ]
            });

        var result = await sut.BuildAsync([], null, new SubmitMessageOptions());

        result.IndexOf("- name: Alpha", StringComparison.Ordinal).Should().BeLessThan(result.IndexOf("- name: Zeta", StringComparison.Ordinal));
    }

    [Fact]
    public async Task BuildAsync_WhenNoToolsExist_UsesFallbackText()
    {
        var sut = CreateBuilder(toolRegistry: new StubToolRegistry());

        var result = await sut.BuildAsync([], null, new SubmitMessageOptions());

        result.Should().Contain("No tools available.");
    }

    [Fact]
    public async Task BuildAsync_WhenNoCommandsExist_UsesFallbackText()
    {
        var sut = CreateBuilder(commandRegistry: new StubCommandRegistry());

        var result = await sut.BuildAsync([], null, new SubmitMessageOptions());

        result.Should().Contain("No slash commands available.");
    }

    [Fact]
    public async Task BuildAsync_WhenMemoryIsEmpty_UsesSelfClosingMemoryElement()
    {
        var sut = CreateBuilder(memoryService: new StubSessionMemoryService { CurrentMemory = "   " });

        var result = await sut.BuildAsync([], null, new SubmitMessageOptions());

        result.Should().Contain("<session_memory />");
    }

    [Fact]
    public async Task BuildAsync_WithBuddyFlag_AddsCompanionSegment()
    {
        var featureFlags = new StubFeatureFlagService();
        featureFlags.EnabledFlags.Add("BUDDY");
        var companionService = new StubCompanionService
        {
            Factory = _ => new Companion(Species.Fox, Eye.Gold, Hat.Crown, Rarity.Legendary, "Nova")
        };
        var sut = CreateBuilder(featureFlags: featureFlags, companionService: companionService);

        var result = await sut.BuildAsync([CreateMessage("hi")], null, new SubmitMessageOptions(Model: "claude"));

        result.Should().Contain("<companion_intro>");
        result.Should().Contain("Companion active: Nova is a Legendary Fox with Gold eyes and a Crown hat.");
        companionService.Seeds.Should().ContainSingle();
    }

    private static SystemPromptBuilder CreateBuilder(
        IToolRegistry? toolRegistry = null,
        ICommandRegistry? commandRegistry = null,
        ISessionMemoryService? memoryService = null,
        IFeatureFlagService? featureFlags = null,
        ICompanionService? companionService = null)
        => new(
            toolRegistry ?? new StubToolRegistry(),
            commandRegistry ?? new StubCommandRegistry(),
            memoryService ?? new StubSessionMemoryService(),
            featureFlags ?? new StubFeatureFlagService(),
            companionService ?? new StubCompanionService(),
            new TestLogger<SystemPromptBuilder>());

    private static Message CreateMessage(string content)
        => new() { MessageId = Guid.NewGuid().ToString("N"), Role = MessageRole.User, Content = content };

    private static JsonElement ParseJson(string json)
    {
        using var document = JsonDocument.Parse(json);
        return document.RootElement.Clone();
    }

    private sealed class StubCommand(string name, string description) : ICommand
    {
        public string Name => name;
        public string[]? Aliases => null;
        public string Description => description;
        public CommandCategory Category => CommandCategory.Local;
        public CommandAvailability Availability => CommandAvailability.Always;
        public bool IsEnabled() => true;
        public Task<CommandResult> ExecuteAsync(CommandContext context, string? args = null, CancellationToken ct = default)
            => Task.FromResult(new CommandResult(true));
    }
}