w4c-workflows-api/w4c-workflows-api.Tests/HttpRequestNodeExecutorQuotaTests.cs

using System.Net;
using System.Text;
using System.Text.Json.Nodes;
using w4c_workflows.Models.Nodes;
using w4c_workflows.Services.Credentials;
using w4c_workflows.Services.Nodes;
using w4c_workflows.Services.Nodes.Executors;
using w4c_workflows.Services.Security;
using Xunit;

namespace w4c_workflows.Tests;

/// <summary>
/// Executor-level quota tests: the response-size cap (declared and chunked
/// bodies) and the run-scoped outbound-request budget.
/// </summary>
public class HttpRequestNodeExecutorQuotaTests
{
    private static readonly NodeBlueprintCatalog Catalog = NodeTestData.CoreCatalog();

    private sealed class Handler : HttpMessageHandler
    {
        private readonly Func<HttpRequestMessage, int, HttpResponseMessage> _responder;

        public Handler(Func<HttpRequestMessage, int, HttpResponseMessage> responder)
            => _responder = responder;

        public List<HttpRequestMessage> Requests { get; } = new();

        protected override Task<HttpResponseMessage> SendAsync(
            HttpRequestMessage request, CancellationToken ct)
        {
            Requests.Add(request);
            return Task.FromResult(_responder(request, Requests.Count - 1));
        }
    }

    private sealed class Factory : IHttpClientFactory
    {
        private readonly HttpMessageHandler _handler;
        public Factory(HttpMessageHandler handler) => _handler = handler;
        public HttpClient CreateClient(string name) => new(_handler, disposeHandler: false);
    }

    /// <summary>Content that deliberately reports no length, forcing the streaming path.</summary>
    private sealed class NoLengthContent : HttpContent
    {
        private readonly byte[] _bytes;
        public NoLengthContent(byte[] bytes) => _bytes = bytes;
        protected override Task SerializeToStreamAsync(Stream stream, TransportContext? context)
            => stream.WriteAsync(_bytes).AsTask();
        protected override bool TryComputeLength(out long length)
        {
            length = 0;
            return false;
        }
    }

    private static HttpResponseMessage Json(string body, HttpStatusCode status = HttpStatusCode.OK)
        => new(status) { Content = new StringContent(body, Encoding.UTF8, "application/json") };

    private static HttpResponseMessage Redirect(string location)
    {
        var response = new HttpResponseMessage(HttpStatusCode.Found)
        {
            Content = new StringContent("{}", Encoding.UTF8, "application/json"),
        };
        response.Headers.TryAddWithoutValidation("Location", location);
        return response;
    }

    private static NodeExecutionContext Context(
        JsonObject parameters, IDictionary<string, object?>? state = null)
        => new()
        {
            Blueprint = Catalog.Get("core.httpRequest")!,
            Parameters = parameters,
            Inputs = new IReadOnlyList<FlowItem>[] { Array.Empty<FlowItem>() },
            State = state ?? new Dictionary<string, object?>(),
        };

    private static HttpRequestNodeExecutor Executor(Handler handler, NodeQuotaPolicy quota)
        => new(new Factory(handler), new CredentialTypeCatalog(), EgressTestData.Guard(), quota, BinaryTestData.Store());

    [Fact]
    public async Task Declared_body_over_the_cap_fails()
    {
        var handler = new Handler((_, _) => Json("""{"data":"too big"}"""));

        var outcome = await Executor(handler, EgressTestData.Quota(o => o.MaxResponseBytes = 4))
            .RunAsync(Context(new JsonObject { ["url"] = "https://example.com/x" }), default);

        Assert.False(outcome.Succeeded);
        Assert.Equal("response_too_large", outcome.Failure!.Code);
    }

    [Fact]
    public async Task Body_exactly_at_the_cap_is_allowed()
    {
        var handler = new Handler((_, _) => new HttpResponseMessage(HttpStatusCode.OK)
        {
            Content = new StringContent("hello", Encoding.UTF8, "text/plain"),
        });

        var outcome = await Executor(handler, EgressTestData.Quota(o => o.MaxResponseBytes = 5))
            .RunAsync(Context(new JsonObject
            {
                ["url"] = "https://example.com/x",
                ["options"] = new JsonObject { ["responseFormat"] = "text" },
            }), default);

        Assert.True(outcome.Succeeded, outcome.Failure?.Message);
        Assert.Equal("hello", Assert.Single(outcome.Outputs[0]).Json["data"]!.GetValue<string>());
    }

    [Fact]
    public async Task Chunked_body_crossing_the_cap_is_rejected_while_streaming()
    {
        var handler = new Handler((_, _) => new HttpResponseMessage(HttpStatusCode.OK)
        {
            Content = new NoLengthContent(Encoding.UTF8.GetBytes(new string('a', 100))),
        });

        var outcome = await Executor(handler, EgressTestData.Quota(o => o.MaxResponseBytes = 10))
            .RunAsync(Context(new JsonObject { ["url"] = "https://example.com/x" }), default);

        Assert.False(outcome.Succeeded);
        Assert.Equal("response_too_large", outcome.Failure!.Code);
    }

    [Fact]
    public async Task Request_budget_exceeded_fails_the_hop()
    {
        var handler = new Handler((_, _) => Redirect("https://example.com/next"));

        var outcome = await Executor(handler, EgressTestData.Quota(o => o.MaxRequestsPerRun = 1))
            .RunAsync(Context(new JsonObject { ["url"] = "https://example.com/x" }), default);

        Assert.False(outcome.Succeeded);
        Assert.Equal("quota_exceeded", outcome.Failure!.Code);
        Assert.Single(handler.Requests);
    }

    [Fact]
    public async Task Request_budget_is_shared_across_nodes_in_one_run()
    {
        var handler = new Handler((_, _) => Json("{}"));
        var executor = Executor(handler, EgressTestData.Quota(o => o.MaxRequestsPerRun = 1));
        var state = new Dictionary<string, object?>();

        var first = await executor.RunAsync(
            Context(new JsonObject { ["url"] = "https://example.com/a" }, state), default);
        var second = await executor.RunAsync(
            Context(new JsonObject { ["url"] = "https://example.com/b" }, state), default);

        Assert.True(first.Succeeded, first.Failure?.Message);
        Assert.False(second.Succeeded);
        Assert.Equal("quota_exceeded", second.Failure!.Code);
    }

    [Fact]
    public async Task Zero_budget_disables_the_limits()
    {
        var handler = new Handler((_, _) => Json("""{"ok":true}"""));
        var quota = EgressTestData.Quota(o =>
        {
            o.MaxRequestsPerRun = 0;
            o.MaxResponseBytes = 0;
        });

        var outcome = await Executor(handler, quota).RunAsync(
            Context(new JsonObject { ["url"] = "https://example.com/x" }), default);

        Assert.True(outcome.Succeeded, outcome.Failure?.Message);
    }
}