internal/wrr/wrr_test.go

/*
 *
 * Copyright 2019 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package wrr

import (
	"errors"
	"math"
	"math/rand"
	"testing"

	"github.com/google/go-cmp/cmp"
	"google.golang.org/grpc/internal/grpctest"
)

type s struct {
	grpctest.Tester
}

func Test(t *testing.T) {
	grpctest.RunSubTests(t, s{})
}

const iterCount = 10000

func equalApproximate(a, b float64) error {
	opt := cmp.Comparer(func(x, y float64) bool {
		delta := math.Abs(x - y)
		mean := math.Abs(x+y) / 2.0
		return delta/mean < 0.05
	})
	if !cmp.Equal(a, b, opt) {
		return errors.New(cmp.Diff(a, b))
	}
	return nil
}

func testWRRNext(t *testing.T, newWRR func() WRR) {
	tests := []struct {
		name    string
		weights []int64
	}{
		{
			name:    "1-1-1",
			weights: []int64{1, 1, 1},
		},
		{
			name:    "1-2-3",
			weights: []int64{1, 2, 3},
		},
		{
			name:    "5-3-2",
			weights: []int64{5, 3, 2},
		},
		{
			name:    "17-23-37",
			weights: []int64{17, 23, 37},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var sumOfWeights int64

			w := newWRR()
			for i, weight := range tt.weights {
				w.Add(i, weight)
				sumOfWeights += weight
			}

			results := make(map[int]int)
			for i := 0; i < iterCount; i++ {
				results[w.Next().(int)]++
			}

			wantRatio := make([]float64, len(tt.weights))
			for i, weight := range tt.weights {
				wantRatio[i] = float64(weight) / float64(sumOfWeights)
			}
			gotRatio := make([]float64, len(tt.weights))
			for i, count := range results {
				gotRatio[i] = float64(count) / iterCount
			}

			for i := range wantRatio {
				if err := equalApproximate(gotRatio[i], wantRatio[i]); err != nil {
					t.Errorf("%v not equal %v", i, err)
				}
			}
		})
	}
}

func (s) TestRandomWRRNext(t *testing.T) {
	testWRRNext(t, NewRandom)
}

func (s) TestEdfWrrNext(t *testing.T) {
	testWRRNext(t, NewEDF)
}

func init() {
	r := rand.New(rand.NewSource(0))
	grpcrandInt63n = r.Int63n
}