package xhttp_test

import (
	"net/http"
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/coder/coder/v2/coderd/util/xhttp"
)

func TestIsRateLimited(t *testing.T) {
	t.Parallel()

	hdr := func(headers map[string]string) http.Header {
		h := http.Header{}
		for k, v := range headers {
			h.Set(k, v)
		}
		return h
	}

	cases := []struct {
		name    string
		status  int
		nilResp bool
		header  map[string]string
		want    bool
	}{
		{name: "Nil", nilResp: true, want: false},
		{name: "OK", status: http.StatusOK, want: false},
		// A successful response with a zeroed remaining count is not a
		// rate-limited rejection.
		{name: "OKZeroRemaining", status: http.StatusOK, header: map[string]string{"X-RateLimit-Remaining": "0"}, want: false},
		{name: "TooManyRequests", status: http.StatusTooManyRequests, want: true},
		{name: "ForbiddenZeroRemaining", status: http.StatusForbidden, header: map[string]string{"X-RateLimit-Remaining": "0"}, want: true},
		{name: "ForbiddenRetryAfter", status: http.StatusForbidden, header: map[string]string{"Retry-After": "60"}, want: true},
		// GitHub secondary limits send Retry-After while the primary quota
		// still has remaining requests; Retry-After alone is sufficient.
		{name: "ForbiddenRetryAfterPositiveRemaining", status: http.StatusForbidden, header: map[string]string{"Retry-After": "60", "X-RateLimit-Remaining": "5000"}, want: true},
		{name: "ForbiddenPositiveRemaining", status: http.StatusForbidden, header: map[string]string{"X-RateLimit-Remaining": "5000"}, want: false},
		// GitLab uses the unprefixed RateLimit-Remaining header.
		{name: "ForbiddenGitLabZeroRemaining", status: http.StatusForbidden, header: map[string]string{"RateLimit-Remaining": "0"}, want: true},
		{name: "ForbiddenGitLabPositiveRemaining", status: http.StatusForbidden, header: map[string]string{"RateLimit-Remaining": "42"}, want: false},
		{name: "ForbiddenNoHeaders", status: http.StatusForbidden, want: false},
		{name: "Unauthorized", status: http.StatusUnauthorized, header: map[string]string{"X-RateLimit-Remaining": "0"}, want: false},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			var resp *http.Response
			if !tc.nilResp {
				resp = &http.Response{StatusCode: tc.status, Header: hdr(tc.header)}
			}
			assert.Equal(t, tc.want, xhttp.IsRateLimited(resp))
		})
	}
}