package ratelimit

import (
	"context"
	"fmt"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/serverless-service/internal/config"
)

func TestLimit(t *testing.T) {
	g := NewGomegaWithT(t)
	r := NewRateLimiter(time.Second, 5)
	c := context.Background()

	for i := range 3 {
		go func(k int) {
			ctx := context.WithValue(c, config.CtxIdentityKey, fmt.Sprintf("test%d@pingcap.com", k))

			for range 5 {
				err := r.Limit(ctx)
				g.Expect(err).To(BeNil())
			}

			for range 2 {
				err := r.Limit(ctx)
				berr, ok := err.(bizerr.BizError)
				g.Expect(ok).To(BeTrue())
				g.Expect(berr.ErrCode()).To(Equal(bizerr.ErrCodeTooManyRequests))
				g.Expect(berr.ErrMsg()).To(Equal("Too many requests"))
			}

			time.Sleep(time.Second)
			for range 5 {
				err := r.Limit(ctx)
				g.Expect(err).To(BeNil())
			}
		}(i)
	}
}
