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

using System.Net;
using System.Text;
using System.Text.Json;
using HSchool.Server.Game;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;

namespace HSchool.Server.Tests;

public class SwarmUiClientTests
{
    [Fact]
    public async Task GenerateAsync_UsesSessionAndReturnsImageBytes()
    {
        var handler = new FakeHandler();
        var http = new HttpClient(handler) { BaseAddress = new Uri("http://swarm.test/") };
        var client = new SwarmUiClient(
            http,
            Options.Create(new SwarmUiOptions { BaseUrl = "http://swarm.test", TimeoutSeconds = 30 }),
            NullLogger<SwarmUiClient>.Instance);

        var settings = new SwarmUiSettings
        {
            Model = "model.safetensors",
            Steps = 8,
            CfgScale = 1,
            Avatar = new SwarmUiSettings.SwarmUiPreset { Width = 512, Height = 512 },
        };

        var bytes = await client.GenerateAsync("a student", "bad", settings, PortraitKind.Avatar, CancellationToken.None);

        Assert.Equal([0x89, 0x50, 0x4E, 0x47], bytes.Take(4));
        Assert.Contains("/API/GetNewSession", handler.Requests[0]);
        Assert.Contains("/API/GenerateText2Image", handler.Requests[1]);
    }

    [Fact]
    public async Task ProbeAsync_ReturnsTrueWhenSessionOpens()
    {
        var handler = new FakeHandler();
        var http = new HttpClient(handler) { BaseAddress = new Uri("http://swarm.test/") };
        var client = new SwarmUiClient(
            http,
            Options.Create(new SwarmUiOptions { BaseUrl = "http://swarm.test", TimeoutSeconds = 30 }),
            NullLogger<SwarmUiClient>.Instance);

        var ok = await client.ProbeAsync(TimeSpan.FromSeconds(5), CancellationToken.None);

        Assert.True(ok);
        Assert.Single(handler.Requests);
        Assert.Contains("/API/GetNewSession", handler.Requests[0]);
    }

    [Fact]
    public async Task ProbeAsync_ReturnsFalseWhenUnreachable()
    {
        var handler = new FailingHandler();
        var http = new HttpClient(handler) { BaseAddress = new Uri("http://swarm.test/") };
        var client = new SwarmUiClient(
            http,
            Options.Create(new SwarmUiOptions { BaseUrl = "http://swarm.test", TimeoutSeconds = 30 }),
            NullLogger<SwarmUiClient>.Instance);

        var ok = await client.ProbeAsync(TimeSpan.FromSeconds(5), CancellationToken.None);

        Assert.False(ok);
    }

    private sealed class FailingHandler : HttpMessageHandler
    {
        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) =>
            Task.FromResult(new HttpResponseMessage(HttpStatusCode.ServiceUnavailable));
    }

    [Fact]
    public async Task GenerateAsync_SendsClipStopAtLayerForClipSkip()
    {
        var handler = new CapturingHandler();
        var http = new HttpClient(handler) { BaseAddress = new Uri("http://swarm.test/") };
        var client = new SwarmUiClient(
            http,
            Options.Create(new SwarmUiOptions { BaseUrl = "http://swarm.test", TimeoutSeconds = 30 }),
            NullLogger<SwarmUiClient>.Instance);

        var settings = new SwarmUiSettings
        {
            Model = "model.safetensors",
            Steps = 4,
            CfgScale = 2,
            ClipSkip = 2,
            Sampler = "dpmpp_sde",
            Scheduler = "karras",
            Avatar = new SwarmUiSettings.SwarmUiPreset { Width = 512, Height = 512 },
        };

        await client.GenerateAsync("a student", "bad", settings, PortraitKind.Avatar, CancellationToken.None);

        using var document = JsonDocument.Parse(handler.GenerateBody!);
        Assert.Equal(-2, document.RootElement.GetProperty("clipstopatlayer").GetInt32());
        Assert.False(document.RootElement.TryGetProperty("clipskip", out _));
    }

    private sealed class CapturingHandler : HttpMessageHandler
    {
        public string? GenerateBody { get; private set; }

        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            if (request.RequestUri!.AbsolutePath.Contains("GetNewSession", StringComparison.Ordinal))
            {
                var session = JsonSerializer.Serialize(new { session_id = "sess-1" });
                return new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(session, Encoding.UTF8, "application/json"),
                };
            }

            if (request.RequestUri.AbsolutePath.Contains("GenerateText2Image", StringComparison.Ordinal))
            {
                GenerateBody = request.Content is null ? null : await request.Content.ReadAsStringAsync(cancellationToken);
                var payload = JsonSerializer.Serialize(new { images = new[] { "data:image/png;base64,iVBORw0KGgo=" } });
                return new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(payload, Encoding.UTF8, "application/json"),
                };
            }

            return new HttpResponseMessage(HttpStatusCode.NotFound);
        }
    }

    private sealed class FakeHandler : HttpMessageHandler
    {
        public List<string> Requests { get; } = [];

        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
        {
            Requests.Add(request.RequestUri!.AbsolutePath);

            if (request.RequestUri.AbsolutePath.Contains("GetNewSession", StringComparison.Ordinal))
            {
                var session = JsonSerializer.Serialize(new { session_id = "sess-1" });
                return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(session, Encoding.UTF8, "application/json"),
                });
            }

            if (request.RequestUri.AbsolutePath.Contains("GenerateText2Image", StringComparison.Ordinal))
            {
                var payload = JsonSerializer.Serialize(new { images = new[] { "data:image/png;base64,iVBORw0KGgo=" } });
                return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK)
                {
                    Content = new StringContent(payload, Encoding.UTF8, "application/json"),
                });
            }

            return Task.FromResult(new HttpResponseMessage(HttpStatusCode.NotFound));
        }
    }
}