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
|
package ratelimit
import (
"context"
"errors"
"fmt"
"strings"
"testing"
)
func TestTokenRateLimit(t *testing.T) {
type usage struct {
Cost uint
Release bool
Err string
AddTokens uint
}
cases := map[string]struct {
Tokens uint
Usages []usage
}{
"retrieve": {
Tokens: 10,
Usages: []usage{
{Cost: 5, Release: true},
{Cost: 5},
{Cost: 5},
{Cost: 5, Err: "retry quota exceeded"},
{AddTokens: 5},
{Cost: 5},
},
},
}
for name, c := range cases {
t.Run(name, func(t *testing.T) {
rl := NewTokenRateLimit(c.Tokens)
for i, u := range c.Usages {
t.Run(fmt.Sprintf("usage_%d", i), func(t *testing.T) {
if u.Cost != 0 {
f, err := rl.GetToken(context.Background(), u.Cost)
if len(u.Err) != 0 {
if err == nil {
t.Fatalf("expect error, got none")
}
if e, a := u.Err, err.Error(); !strings.Contains(a, e) {
t.Fatalf("expect %q error, got %q", e, a)
}
} else if err != nil {
t.Fatalf("expect no error, got %v", err)
}
if u.Release {
if err := f(); err != nil {
t.Fatalf("expect no error, got %v", err)
}
}
}
if u.AddTokens != 0 {
rl.AddTokens(u.AddTokens)
}
})
}
})
}
}
func TestTokenRateLimit_canceled(t *testing.T) {
rl := NewTokenRateLimit(10)
ctx, cancel := context.WithCancel(context.Background())
cancel()
fn, err := rl.GetToken(ctx, 1)
if err == nil {
t.Fatalf("expect error, got none")
}
if fn != nil {
t.Errorf("expect no release func returned")
}
var v interface{ CanceledError() bool }
if !errors.As(err, &v) {
t.Fatalf("expect %T error, got %v", v, err)
}
}
|
