/
testdata_test.go
87 lines (75 loc) · 2.14 KB
/
testdata_test.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
// Copyright (C) 2015 Foursquare Labs Inc.
package main
import (
"math/rand"
"sort"
"github.com/foursquare/quiver/gen"
"github.com/foursquare/quiver/hfile"
"github.com/stretchr/testify/assert"
)
var uncompressed *ThriftRpcImpl
var compressed *ThriftRpcImpl
var compressedMapped *ThriftRpcImpl
var maxKey int
type hasFatal interface {
Fatal(args ...interface{})
Fatalf(format string, args ...interface{})
assert.TestingT
}
func Setup(t hasFatal) {
maxKey = 10000000
if uncompressed != nil && compressed != nil {
return
}
if cs, err := hfile.TestdataCollectionSet("uncompressed", maxKey, false, hfile.CopiedToMem); err != nil {
t.Fatal(err)
} else {
uncompressed = &ThriftRpcImpl{&RpcShared{cs}}
}
if cs, err := hfile.TestdataCollectionSet("compressed", maxKey, true, hfile.CopiedToMem); err != nil {
t.Fatal(err)
} else {
compressed = &ThriftRpcImpl{&RpcShared{cs}}
}
}
func SetupMapped(t hasFatal) {
if cs, err := hfile.TestdataCollectionSet("compressed", maxKey, true, hfile.MemlockFile); err != nil {
t.Fatal(err)
} else {
compressedMapped = &ThriftRpcImpl{&RpcShared{cs}}
}
}
func GetTestIntReq(name string, keys []int) *gen.SingleHFileKeyRequest {
keyBytes := make([][]byte, len(keys))
for i, v := range keys {
keyBytes[i] = hfile.MockKeyInt(v)
}
return &gen.SingleHFileKeyRequest{&name, keyBytes, nil, nil}
}
func MakeTestKeyIntList(r *rand.Rand, count, max int) []int {
if count < 1 {
return nil
}
keys := make([]int, count)
for i := 0; i < count; i++ {
keys[i] = r.Intn(max)
}
sort.Ints(keys)
return keys
}
func GetRandomTestReqs(name string, count, reqSize, max int) []*gen.SingleHFileKeyRequest {
r := rand.New(rand.NewSource(int64(count + reqSize + max)))
reqs := make([]*gen.SingleHFileKeyRequest, count)
for i := 0; i < count; i++ {
reqs[i] = GetTestIntReq(name, MakeTestKeyIntList(r, reqSize, max))
}
return reqs
}
func CheckReqAndRes(t assert.TestingT, req *gen.SingleHFileKeyRequest, res *gen.SingleHFileKeyResponse) {
for i, k := range req.GetSortedKeys() {
expected := hfile.MockValueForMockKey(k)
actual, ok := res.Values[int32(i)]
assert.True(t, ok)
assert.Equal(t, actual, expected)
}
}