package tidbmanager

import (
	"context"
	"testing"
	"time"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"golang.org/x/time/rate"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func Test_rateLimiter_wait(t1 *testing.T) {
	tidbPoolNamespace := "tidb-pool"
	tests := []struct {
		name           string
		poolConfigData map[string]string
		rateLimit      int
		limiter        *rate.Limiter
		waitTimes      int
		wantMinElapsed time.Duration
		wantRateLimit  int
		wantLimiterNil bool
	}{
		{
			name:           "rate limit works",
			poolConfigData: map[string]string{types.AssignRateLimitPerManager: "2"},
			waitTimes:      2,
			wantMinElapsed: 400 * time.Millisecond,
			wantRateLimit:  2,
			wantLimiterNil: false,
		},
		{
			name:           "rate limit can be updated by configmap",
			poolConfigData: map[string]string{types.AssignRateLimitPerManager: "0"},
			rateLimit:      2,
			limiter:        rate.NewLimiter(rate.Limit(2), rateLimiterBurst),
			waitTimes:      1,
			wantRateLimit:  0,
			wantLimiterNil: true,
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)
			ctx := context.Background()
			configMap := &corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.ServerlessPoolConfig,
					Namespace: tidbPoolNamespace,
				},
				Data: tt.poolConfigData,
			}
			kubeClient := fake.NewClientBuilder().
				WithScheme(scheme.Scheme).
				WithRuntimeObjects(configMap).
				Build()
			limiter := newRateLimiter(logr.New(log.NullLogSink{}), tidbPoolNamespace, kubeClient)
			limiter.rateLimit = tt.rateLimit
			limiter.limiter = tt.limiter

			start := time.Now()
			for i := 0; i < tt.waitTimes; i++ {
				g.Expect(limiter.wait(ctx)).Should(Succeed())
			}
			if tt.wantMinElapsed > 0 {
				g.Expect(time.Since(start)).Should(BeNumerically(">=", tt.wantMinElapsed))
			}
			g.Expect(limiter.rateLimit).Should(Equal(tt.wantRateLimit))
			if tt.wantLimiterNil {
				g.Expect(limiter.limiter).Should(BeNil())
			} else {
				g.Expect(limiter.limiter).ShouldNot(BeNil())
			}
		})
	}
}