package cluster import ( "context" "encoding/json" "fmt" "net/http" "net/http/httptest" "net/url" "testing" "time" "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1" "github.com/stretchr/testify/require" cloudv1 "github.com/tidbcloud/infra-api/api/v1" "github.com/tidbcloud/shared-provider/pkg/scheme" "github.com/tidbcloud/shared-provider/pkg/types" apitypes "k8s.io/apimachinery/pkg/types" "k8s.io/apimachinery/pkg/util/clock" "sigs.k8s.io/controller-runtime/pkg/client/fake" ) var emptyResult = ` { "status": "success", "data": { "resultType": "vector", "result": [ ] } } ` var responseTPL = ` { "status": "success", "data": { "resultType": "vector", "result": [ { "metric": { "type": "write" }, "value": [ 1629183451.08, "%f" ] }, { "metric": { "type": "read" }, "value": [ 1629183451.08, "%f" ] } ] } } ` func responseWithReadWrite(read float64, write float64) string { if write == 0 && read == 0 { return emptyResult } return fmt.Sprintf(responseTPL, write, read) } func TestTimePeriod(t *testing.T) { now := time.Date(2000, 1, 1, 1, 1, 1, 1, time.Local) for i := 0; i < 7; i++ { now = now.AddDate(0, 0, i) start, end := timePeriod(now) t.Log(now, start, end) require.True(t, now.After(start)) require.True(t, now.Before(end)) require.Equal(t, start.Weekday(), time.Monday) require.Equal(t, start.Hour(), 0) require.Equal(t, start.Minute(), 0) require.Equal(t, start.Second(), 0) require.Equal(t, end.Sub(start), time.Hour*24*7) } } func TestBandwidth(t *testing.T) { var readBytes float64 var writeBytes float64 var lastQuery url.Values ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { lastQuery = r.URL.Query() fmt.Fprintln(w, responseWithReadWrite(readBytes, writeBytes)) })) defer ts.Close() querier := newHTTPClientQueyer(&http.Client{}, ts.URL) key := apitypes.NamespacedName{Namespace: "ns", Name: "name"} tc := new(v1alpha1.TidbCluster) tc.Namespace = key.Namespace tc.Name = key.Name cc := new(cloudv1.Cluster) cc.Namespace = key.Namespace cc.Name = key.Name client := fake.NewClientBuilder(). WithScheme(scheme.Scheme). WithRuntimeObjects(tc). Build() infraClient := fake.NewClientBuilder(). WithScheme(scheme.Scheme). WithRuntimeObjects(cc). Build() ctx := context.Background() h := NewBandwidthHandler(client, infraClient) assertThrottl := func(t *testing.T) { t.Helper() getcc := new(cloudv1.Cluster) err := infraClient.Get(context.Background(), key, getcc) require.Nil(t, err) gettc := new(v1alpha1.TidbCluster) err = client.Get(context.Background(), key, gettc) require.Nil(t, err) v, ok := getcc.Annotations[types.AnnoKeyIngress] v2, ok2 := gettc.Annotations[types.AnnoKeyIngress] require.Equal(t, ok, ok2) require.Equal(t, v, v2) v2, ok2 = gettc.Annotations[types.AnnoKeyIngressPort] require.Equal(t, ok, ok2) if ok { require.Equal(t, v2, "4000") } v, ok = getcc.Annotations[types.AnnoKeyEgress] v2, ok2 = gettc.Annotations[types.AnnoKeyEgress] require.Equal(t, ok, ok2) require.Equal(t, v, v2) v2, ok2 = gettc.Annotations[types.AnnoKeyEgressPort] require.Equal(t, ok, ok2) if ok { require.Equal(t, v2, "4000") } } getTrafficUsage := func(t *testing.T) (trafficUsage *TrafficUsageInfo) { t.Helper() data, ok := cc.Annotations[trafficUsageAnnKey] require.True(t, ok) trafficUsage = new(TrafficUsageInfo) err := json.Unmarshal([]byte(data), trafficUsage) require.Nil(t, err) return } now := time.Date(2020, 8, 16, 0, 0, 0, 0, time.UTC) // monday 00:00 // test syncAnnToTC h.clock = clock.NewFakeClock(now) _, err := h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) // set ann and check synced to tc cc.Annotations = map[string]string{ "other": "v", types.AnnoKeyIngress: "1Gi/s", types.AnnoKeyEgress: "1Gi/s", } err = infraClient.Update(ctx, cc) require.Nil(t, err) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) require.Equal(t, cc.Annotations["other"], "v") // clear and check synced to tc delete(cc.Annotations, types.AnnoKeyIngress) delete(cc.Annotations, types.AnnoKeyEgress) err = infraClient.Update(ctx, cc) require.Nil(t, err) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) require.Equal(t, cc.Annotations["other"], "v") now = time.Date(2021, 8, 16, 0, 0, 0, 0, time.UTC) // monday 00:00 require.Equal(t, time.Monday, now.Weekday()) // test bandwidth info key readBytes = 1024 writeBytes = 1024 h.clock = clock.NewFakeClock(now) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) trafficUsage := getTrafficUsage(t) require.Equal(t, trafficUsage.IngressUsage, 1024.0) require.Equal(t, trafficUsage.EgressUsage, 1024.0, now) require.Equal(t, now.String(), trafficUsage.LastAccumulateTime.String()) timeParam, err := time.Parse(time.RFC3339, lastQuery.Get("time")) require.Nil(t, err) require.Equal(t, timeParam.String(), now.String()) // accumulate more bytes now = now.AddDate(0, 0, 1) readBytes = float64(TotalInBytes) writeBytes = float64(TotalInBytes) + 1024.0 h.clock = clock.NewFakeClock(now) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) trafficUsage = getTrafficUsage(t) require.Equal(t, trafficUsage.IngressUsage, readBytes+1024 /*1024 for the first time*/) require.Equal(t, trafficUsage.EgressUsage, writeBytes+1024 /*1024 for the first time*/) require.Equal(t, trafficUsage.LastAccumulateTime, now) // cross to a new period will clear accumulated bytes by last period now = now.AddDate(0, 0, 7) readBytes = 1024 writeBytes = 2046 h.clock = clock.NewFakeClock(now) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) trafficUsage = getTrafficUsage(t) require.Equal(t, trafficUsage.IngressUsage, readBytes) require.Equal(t, trafficUsage.EgressUsage, writeBytes) require.Equal(t, trafficUsage.LastAccumulateTime, now) // should still update LastAccumulateTime if no metric now = now.AddDate(0, 0, 1) readBytes = 0 writeBytes = 0 h.clock = clock.NewFakeClock(now) _, err = h.Sync(ctx, cc, tc, querier) require.Nil(t, err) assertThrottl(t) trafficUsage = getTrafficUsage(t) require.Equal(t, trafficUsage.LastAccumulateTime, now) // check returning checkAfter as expect pass := time.Minute now = now.Add(time.Minute) h.clock = clock.NewFakeClock(now) checkAfter, err := h.Sync(ctx, cc, tc, querier) require.Nil(t, err) require.Equal(t, checkAfter, defaultAccumulatePeriodDuration-pass) }