h-school/tests/HSchool.Server.Tests/SwarmUiSettingsStoreTests.cs

using HSchool.Server.Game;

namespace HSchool.Server.Tests;

public class SwarmUiSettingsStoreTests
{
    [Fact]
    public void Resolve_UsesAgeRulePreset()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var profile = config.Resolve(8, PortraitKind.Avatar);

        Assert.Equal("child", profile.PresetId);
        Assert.Equal(SwarmUiModelDefinition.DreamShaperId, profile.Model);
        Assert.Equal(4, profile.Steps);
        Assert.Equal(2, profile.CfgScale);
    }

    [Fact]
    public void Resolve_FallsBackToActivePresetForAdults()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var profile = config.Resolve(34, PortraitKind.Full);

        Assert.Equal("default", profile.PresetId);
        Assert.Equal(SwarmUiModelDefinition.BabesId, profile.Model);
        Assert.Equal(7, profile.Steps);
        Assert.Equal(1.5, profile.CfgScale);
    }

    [Fact]
    public void Resolve_WithoutOverride_TakesStepsAndCfgFromModel()
    {
        var config = new SwarmUiConfigFile
        {
            Models =
            [
                new SwarmUiModelDefinition { Id = "m.safetensors", Steps = 9, CfgScale = 2.25 },
            ],
            Presets =
            [
                new SwarmUiPresetDefinition
                {
                    Id = "default",
                    Label = "Default",
                    Model = "m.safetensors",
                    Avatar = new SwarmUiKindPreset { Width = 512, Height = 512 },
                    Custom = new SwarmUiKindPreset { Width = 512, Height = 512 },
                    FullBody = new SwarmUiKindPreset { Width = 512, Height = 512 },
                },
            ],
        };

        var profile = config.Resolve(20, PortraitKind.Avatar);

        Assert.Equal(9, profile.Steps);
        Assert.Equal(2.25, profile.CfgScale);
    }

    [Fact]
    public void ToProfile_JoinsModelAndPresetNegative()
    {
        var model = new SwarmUiModelDefinition { Id = "m.safetensors", Negative = "model-neg" };
        var preset = new SwarmUiPresetDefinition
        {
            Id = "default",
            Label = "Default",
            Model = "m.safetensors",
            Negative = "preset-neg",
            Avatar = new SwarmUiKindPreset { Width = 512, Height = 512 },
        };

        var profile = preset.ToProfile(PortraitKind.Avatar, model);

        Assert.Equal("model-neg, preset-neg", profile.Negative);
    }

    [Fact]
    public void Resolve_PresetOverrideBeatsModelDefaults()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var preset = config.Presets[0];
        preset.Steps = 12;
        preset.CfgScale = 3.5;

        var profile = config.Resolve(34, PortraitKind.Avatar);

        Assert.Equal(12, profile.Steps);
        Assert.Equal(3.5, profile.CfgScale);
        Assert.Equal(0, profile.ClipSkip);
    }

    [Fact]
    public void NormalizeAfterLoad_MigratesLegacyFlatFile()
    {
        var config = new SwarmUiConfigFile
        {
            Model = "legacy.safetensors",
            Steps = 12,
            Positive = "legacy prompt",
            Avatar = new SwarmUiKindPreset { Width = 640, Height = 640, Positive = "close up." },
        };

        config.NormalizeAfterLoad();

        Assert.Single(config.Presets);
        Assert.Equal("default", config.Presets[0].Id);
        Assert.Equal("legacy.safetensors", config.Presets[0].Model);
        Assert.Equal("legacy prompt", config.Presets[0].Style);
        Assert.Null(config.Presets[0].Positive);
        Assert.Null(config.Presets[0].Steps);
        var model = Assert.Single(config.Models, entry => entry.Id == "legacy.safetensors");
        Assert.Equal(12, model.Steps);
        Assert.Equal("close up.", config.Presets[0].Avatar!.ShotType);
        Assert.Null(config.Presets[0].Avatar!.Positive);
    }

    [Fact]
    public void NormalizeAfterLoad_LiftsLegacyPresetPositiveIntoStyleAndShotType()
    {
        var config = new SwarmUiConfigFile
        {
            Presets =
            [
                new SwarmUiPresetDefinition
                {
                    Id = "default",
                    Label = "Default",
                    Model = "a.safetensors",
                    Steps = 6,
                    CfgScale = 2,
                    Positive = "was base prompt",
                    Negative = "old neg",
                    Avatar = new SwarmUiKindPreset { Width = 512, Height = 512, Positive = "was kind" },
                    Custom = new SwarmUiKindPreset { Width = 512, Height = 512 },
                    FullBody = new SwarmUiKindPreset { Width = 512, Height = 768 },
                },
            ],
        };

        config.NormalizeAfterLoad();

        Assert.Equal("was base prompt", config.Presets[0].Style);
        Assert.Equal("was kind", config.Presets[0].Avatar!.ShotType);
        Assert.Contains(config.Models, model => model.Id == "a.safetensors");
        Assert.Equal("", Assert.Single(config.Models).Positive);
    }

    [Fact]
    public void AllowedModelIds_IntersectsSwarmWithCatalog()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var discovery = new SwarmUiDiscovery(
            true,
            [
                SwarmUiModelDefinition.BabesId,
                "not-in-config.safetensors",
                SwarmUiModelDefinition.EpicRealismId,
            ],
            [],
            [],
            [],
            []);

        var allowed = config.AllowedModelIds(discovery);

        Assert.Equal(
            [SwarmUiModelDefinition.BabesId, SwarmUiModelDefinition.EpicRealismId],
            allowed);
        Assert.DoesNotContain("not-in-config.safetensors", allowed);
        Assert.DoesNotContain(SwarmUiModelDefinition.DreamShaperId, allowed);
    }

    [Fact]
    public void AllowedModelIds_WhenSwarmOffline_ReturnsCatalog()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var allowed = config.AllowedModelIds(SwarmUiDiscovery.Offline);

        Assert.Equal(config.Models.Select(model => model.Id), allowed);
    }

    [Fact]
    public void Clone_IsIndependentOfTheSource()
    {
        var source = SwarmUiConfigFile.CreateDefault();
        source.Presets[0].Model = "mutated.safetensors";
        source.Models[0].Id = "mutated.safetensors";

        var copy = SwarmUiConfigFile.Clone(source);
        copy.Presets[0].Model = "other.safetensors";

        Assert.Equal("mutated.safetensors", source.Presets[0].Model);
        Assert.Equal("other.safetensors", copy.Presets[0].Model);
        Assert.Equal(source.ActivePresetId, copy.ActivePresetId);
    }

    [Fact]
    public void CreateDefault_SeedsEveryKnownSwarmModel()
    {
        var config = SwarmUiConfigFile.CreateDefault();
        var ids = config.Models.Select(model => model.Id).ToList();

        Assert.Contains(SwarmUiModelDefinition.BabesId, ids);
        Assert.Contains(SwarmUiModelDefinition.DreamShaperId, ids);
        Assert.Contains(SwarmUiModelDefinition.EpicRealismId, ids);
        Assert.Contains(SwarmUiModelDefinition.LustifyId, ids);
    }
}