forked from dolthub/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
throttlerclient_testsuite.go
140 lines (113 loc) · 3.88 KB
/
throttlerclient_testsuite.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// Copyright 2016, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package throttlerclienttest contains the testsuite against which each
// RPC implementation of the throttlerclient interface must be tested.
package throttlerclienttest
// NOTE: This file is not test-only code because it is referenced by
// tests in other packages and therefore it has to be regularly
// visible.
// NOTE: This code is in its own package such that its dependencies
// (e.g. zookeeper) won't be drawn into production binaries as well.
import (
"reflect"
"strings"
"testing"
"golang.org/x/net/context"
"github.com/youtube/vitess/go/vt/throttler"
"github.com/youtube/vitess/go/vt/throttler/throttlerclient"
)
// TestSuite runs the test suite on the given throttlerclient and throttlerserver.
func TestSuite(t *testing.T, c throttlerclient.Client) {
tf := &testFixture{}
if err := tf.setUp(); err != nil {
t.Fatal(err)
}
defer tf.tearDown()
tf.maxRates(t, c)
tf.setMaxRate(t, c)
}
// TestSuitePanics tests the panic handling of each RPC method. Unlike TestSuite
// it does not use the real throttler.managerImpl. Instead, it uses FakeManager
// which allows us to panic on each RPC.
func TestSuitePanics(t *testing.T, c throttlerclient.Client) {
maxRatesPanics(t, c)
setMaxRatePanics(t, c)
}
var throttlerNames = []string{"t1", "t2"}
type testFixture struct {
throttlers []*throttler.Throttler
}
func (tf *testFixture) setUp() error {
for _, name := range throttlerNames {
t, err := throttler.NewThrottler(name, "TPS", 1 /* threadCount */, 1, throttler.ReplicationLagModuleDisabled)
if err != nil {
return err
}
tf.throttlers = append(tf.throttlers, t)
}
return nil
}
func (tf *testFixture) tearDown() {
for _, t := range tf.throttlers {
t.Close()
}
}
func (tf *testFixture) maxRates(t *testing.T, client throttlerclient.Client) {
_, err := client.SetMaxRate(context.Background(), 23)
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
got, err := client.MaxRates(context.Background())
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
want := map[string]int64{
"t1": 23,
"t2": 23,
}
if !reflect.DeepEqual(got, want) {
t.Fatalf("rate was not updated on all registered throttlers. got = %v, want = %v", got, throttlerNames)
}
}
func (tf *testFixture) setMaxRate(t *testing.T, client throttlerclient.Client) {
got, err := client.SetMaxRate(context.Background(), 23)
if err != nil {
t.Fatalf("Cannot execute remote command: %v", err)
}
if !reflect.DeepEqual(got, throttlerNames) {
t.Fatalf("rate was not updated on all registered throttlers. got = %v, want = %v", got, throttlerNames)
}
}
// FakeManager implements the throttler.Manager interface and panics on all
// methods defined in the interface.
type FakeManager struct {
}
const panicMsg = "RPC server implementation should handle this"
// MaxRates implements the throttler.Manager interface. It always panics.
func (fm *FakeManager) MaxRates() map[string]int64 {
panic(panicMsg)
}
// SetMaxRate implements the throttler.Manager interface. It always panics.
func (fm *FakeManager) SetMaxRate(int64) []string {
panic(panicMsg)
}
// Test methods which test for each RPC that panics are caught.
func maxRatesPanics(t *testing.T, client throttlerclient.Client) {
_, err := client.MaxRates(context.Background())
if !errorFromPanicHandler(err) {
t.Fatalf("MaxRates RPC implementation does not catch panics properly: %v", err)
}
}
func setMaxRatePanics(t *testing.T, client throttlerclient.Client) {
_, err := client.SetMaxRate(context.Background(), 23)
if !errorFromPanicHandler(err) {
t.Fatalf("SetMaxRate RPC implementation does not catch panics properly: %v", err)
}
}
func errorFromPanicHandler(err error) bool {
if err == nil || !strings.Contains(err.Error(), panicMsg) {
return false
}
return true
}