asconnect-service/Tests/ASConnectService/Cases/Public/Middlewares/RateLimitMiddlewareTests.swift

import Foundation
import HTTPTypes
import OpenAPIRuntime
import Testing

// A testable import gives access to the initializer that injects the sleep closure.
@testable import ASConnectService

@Suite("RateLimitMiddleware")
struct RateLimitMiddlewareTests {
    // MARK: Properties

    private let recorder = Recorder()

    // MARK: Tests

    @Test("Returns a response that is not rate limited without retrying")
    func returnsResponseWithoutRetrying() async throws {
        let response = try await intercept(responses: [.ok])

        #expect(response.status == .ok)
        #expect(await recorder.calls == 1)
        #expect(await recorder.delays.isEmpty)
    }

    @Test("Retries rate-limited requests with an exponential backoff")
    func retriesWithBackoff() async throws {
        let response = try await intercept(responses: [.tooManyRequests, .tooManyRequests, .ok])

        #expect(response.status == .ok)
        #expect(await recorder.calls == 3)
        #expect(await recorder.delays == [1, 2])
    }

    @Test("Waits for the seconds in the Retry-After header")
    func waitsForRetryAfter() async throws {
        let response = try await intercept(responses: [.tooManyRequests(retryAfter: "5"), .ok])

        #expect(response.status == .ok)
        #expect(await recorder.delays == [5])
    }

    @Test("Returns the rate-limited response when the retries run out")
    func returnsResponseWhenRetriesRunOut() async throws {
        let response = try await intercept(
            responses: [.tooManyRequests, .tooManyRequests, .tooManyRequests],
            maximumRetries: 2
        )

        #expect(response.status == .tooManyRequests)
        #expect(await recorder.calls == 3)
        #expect(await recorder.delays == [1, 2])
    }

    @Test("Returns the rate-limited response when the delay exceeds the maximum")
    func returnsResponseWhenDelayExceedsMaximum() async throws {
        let response = try await intercept(responses: [.tooManyRequests(retryAfter: "3600"), .ok])

        #expect(response.status == .tooManyRequests)
        #expect(await recorder.calls == 1)
        #expect(await recorder.delays.isEmpty)
    }

    @Test("Does not retry a request whose body can only be sent once")
    func doesNotRetrySingleIterationBody() async throws {
        let body = HTTPBody(
            AsyncStream<HTTPBody.ByteChunk> { $0.finish() },
            length: .unknown,
            iterationBehavior: .single
        )

        let response = try await intercept(responses: [.tooManyRequests, .ok], body: body)

        #expect(response.status == .tooManyRequests)
        #expect(await recorder.calls == 1)
    }
}

// MARK: - Helpers

private extension RateLimitMiddlewareTests {
    func intercept(
        responses: [HTTPResponse],
        body: HTTPBody? = nil,
        maximumRetries: Int = 3
    ) async throws -> HTTPResponse {
        let recorder = recorder
        let middleware = RateLimitMiddleware(
            maximumRetries: maximumRetries,
            baseDelay: 1,
            maximumDelay: 60,
            sleep: { await recorder.record(delay: $0) }
        )

        let (response, _) = try await middleware.intercept(
            .get,
            body: body,
            baseURL: .base,
            operationID: "apps_getCollection",
            next: { _, _, _ in
                (await recorder.response(from: responses), nil)
            }
        )

        return response
    }
}

/// An actor that records the calls to the next handler, and the delays the middleware waits for.
private actor Recorder {
    private(set) var calls = 0
    private(set) var delays: [TimeInterval] = []

    func response(from responses: [HTTPResponse]) -> HTTPResponse {
        defer { calls += 1 }

        return responses[calls]
    }

    func record(delay: TimeInterval) {
        delays.append(delay)
    }
}

// MARK: - HTTPRequest+Samples

private extension HTTPRequest {
    static let get = HTTPRequest(
        method: .get,
        scheme: "https",
        authority: "api.appstoreconnect.apple.com",
        path: "/v1/apps"
    )
}

// MARK: - HTTPResponse+Samples

private extension HTTPResponse {
    static let ok = HTTPResponse(status: .ok)
    static let tooManyRequests = HTTPResponse(status: .tooManyRequests)

    static func tooManyRequests(retryAfter: String) -> HTTPResponse {
        HTTPResponse(status: .tooManyRequests, headerFields: [.retryAfter: retryAfter])
    }
}

// MARK: - URL+Samples

private extension URL {
    static let base = URL(string: "https://api.appstoreconnect.apple.com")!
}