github.com/google/martian/v3@v3.3.3/trafficshape/bucket_test.go

// Copyright 2015 Google Inc. All rights reserved.
//
// 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 trafficshape

import (
	"errors"
	"runtime"
	"sync/atomic"
	"testing"
	"time"
)

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

	b := NewBucket(10, 10*time.Millisecond)
	defer b.Close()

	if got, want := b.Capacity(), int64(10); got != want {
		t.Fatalf("b.Capacity(): got %d, want %d", got, want)
	}

	n, err := b.Fill(func(remaining int64) (int64, error) {
		if want := int64(10); remaining != want {
			t.Errorf("remaining: got %d, want %d", remaining, want)
		}
		return 5, nil
	})
	if err != nil {
		t.Fatalf("Fill(): got %v, want no error", err)
	}
	if got, want := n, int64(5); got != want {
		t.Fatalf("n: got %d, want %d", got, want)
	}

	n, err = b.Fill(func(remaining int64) (int64, error) {
		if want := int64(5); remaining != want {
			t.Errorf("remaining: got %d, want %d", remaining, want)
		}
		return 5, nil
	})
	if err != nil {
		t.Fatalf("Fill(): got %v, want no error", err)
	}
	if got, want := n, int64(5); got != want {
		t.Fatalf("n: got %d, want %d", got, want)
	}
	n, err = b.Fill(func(remaining int64) (int64, error) {
		t.Fatal("Fill: executed func when full, want skipped")
		return 0, nil
	})
	if err != nil {
		t.Fatalf("Fill(): got %v, want no error", err)
	}

	// Wait for the bucket to drain.
	for {
		if atomic.LoadInt64(&b.fill) == 0 {
			break
		}
		// Allow for a goroutine switch, required for GOMAXPROCS = 1.
		runtime.Gosched()
	}

	wanterr := errors.New("fill function error")
	n, err = b.Fill(func(remaining int64) (int64, error) {
		if want := int64(10); remaining != want {
			t.Errorf("remaining: got %d, want %d", remaining, want)
		}
		return 0, wanterr
	})
	if err != wanterr {
		t.Fatalf("Fill(): got %v, want %v", err, wanterr)
	}
	if got, want := n, int64(0); got != want {
		t.Fatalf("n: got %d, want %d", got, want)
	}
}

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

	b := NewBucket(0, time.Millisecond)
	b.Close()

	if _, err := b.Fill(nil); err != errFillClosedBucket {
		t.Errorf("Fill(): got %v, want errFillClosedBucket", err)
	}
	if _, err := b.FillThrottle(nil); err != errFillClosedBucket {
		t.Errorf("FillThrottle(): got %v, want errFillClosedBucket", err)
	}
}

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

	b := NewBucket(10, 10*time.Millisecond)
	defer b.Close()

	n, err := b.Fill(func(remaining int64) (int64, error) {
		return 11, nil
	})
	if err != nil {
		t.Fatalf("Fill(): got %v, want no error", err)
	}

	n, err = b.Fill(func(int64) (int64, error) {
		t.Fatal("Fill: executed func when full, want skipped")
		return 0, nil
	})
	if err != ErrBucketOverflow {
		t.Fatalf("Fill(): got %v, want ErrBucketOverflow", err)
	}
	if got, want := n, int64(0); got != want {
		t.Fatalf("n: got %d, want %d", got, want)
	}
}

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

	b := NewBucket(50, 50*time.Millisecond)
	defer b.Close()

	closec := make(chan struct{})
	errc := make(chan error, 1)

	fill := func() {
		for {
			select {
			case <-closec:
				return
			default:
				if _, err := b.FillThrottle(func(remaining int64) (int64, error) {
					if remaining < 10 {
						return remaining, nil
					}
					return 10, nil
				}); err != nil {
					select {
					case errc <- err:
					default:
					}
				}
			}
		}
	}

	for i := 0; i < 5; i++ {
		go fill()
	}

	time.Sleep(time.Second)

	close(closec)

	select {
	case err := <-errc:
		t.Fatalf("FillThrottle: got %v, want no error", err)
	default:
	}
}

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

	b := NewBucket(0, time.Minute)
	time.AfterFunc(time.Second, func() { b.Close() })

	if _, err := b.FillThrottle(func(int64) (int64, error) {
		t.Fatal("FillThrottle(): executed func after close, want skipped")
		return 0, nil
	}); err != errFillClosedBucket {
		t.Errorf("b.FillThrottle(): got nil, want errFillClosedBucket")
	}
}