forked from grafana/loki
/
fixtures.go
119 lines (100 loc) · 2.82 KB
/
fixtures.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
package gcp
import (
"context"
"fmt"
"io"
"cloud.google.com/go/bigtable"
"cloud.google.com/go/bigtable/bttest"
"cloud.google.com/go/storage"
"github.com/fsouza/fake-gcs-server/fakestorage"
"google.golang.org/api/option"
"google.golang.org/grpc"
"github.com/frelon/loki/v2/pkg/storage/chunk"
"github.com/frelon/loki/v2/pkg/storage/chunk/hedging"
"github.com/frelon/loki/v2/pkg/storage/chunk/objectclient"
"github.com/frelon/loki/v2/pkg/storage/chunk/testutils"
)
const (
proj, instance = "proj", "instance"
)
type fixture struct {
btsrv *bttest.Server
gcssrv *fakestorage.Server
name string
gcsObjectClient bool
columnKeyClient bool
hashPrefix bool
}
func (f *fixture) Name() string {
return f.name
}
func (f *fixture) Clients() (
iClient chunk.IndexClient, cClient chunk.Client, tClient chunk.TableClient,
schemaConfig chunk.SchemaConfig, closer io.Closer, err error,
) {
f.btsrv, err = bttest.NewServer("localhost:0")
if err != nil {
return
}
f.gcssrv = fakestorage.NewServer(nil)
f.gcssrv.CreateBucket("chunks")
conn, err := grpc.Dial(f.btsrv.Addr, grpc.WithInsecure())
if err != nil {
return
}
ctx := context.Background()
adminClient, err := bigtable.NewAdminClient(ctx, proj, instance, option.WithGRPCConn(conn))
if err != nil {
return
}
schemaConfig = testutils.DefaultSchemaConfig("gcp-columnkey")
tClient = &tableClient{
client: adminClient,
}
client, err := bigtable.NewClient(ctx, proj, instance, option.WithGRPCConn(conn))
if err != nil {
return
}
cfg := Config{
DistributeKeys: f.hashPrefix,
}
if f.columnKeyClient {
iClient = newStorageClientColumnKey(cfg, schemaConfig, client)
} else {
iClient = newStorageClientV1(cfg, schemaConfig, client)
}
if f.gcsObjectClient {
var c *GCSObjectClient
c, err = newGCSObjectClient(ctx, GCSConfig{BucketName: "chunks"}, hedging.Config{}, func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) {
return f.gcssrv.Client(), nil
})
if err != nil {
return
}
cClient = objectclient.NewClient(c, nil, chunk.SchemaConfig{})
} else {
cClient = newBigtableObjectClient(Config{}, schemaConfig, client)
}
closer = testutils.CloserFunc(func() error {
conn.Close()
return nil
})
return
}
// Fixtures for unit testing GCP storage.
var Fixtures = func() []testutils.Fixture {
fixtures := []testutils.Fixture{}
for _, gcsObjectClient := range []bool{true, false} {
for _, columnKeyClient := range []bool{true, false} {
for _, hashPrefix := range []bool{true, false} {
fixtures = append(fixtures, &fixture{
name: fmt.Sprintf("bigtable-columnkey:%v-gcsObjectClient:%v-hashPrefix:%v", columnKeyClient, gcsObjectClient, hashPrefix),
columnKeyClient: columnKeyClient,
gcsObjectClient: gcsObjectClient,
hashPrefix: hashPrefix,
})
}
}
}
return fixtures
}()