forked from vippsas/go-cosmosdb
/
cosmostest.go
146 lines (134 loc) · 4.37 KB
/
cosmostest.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
141
142
143
144
145
146
// The cosmostest package contains utilities for writing tests with cosmos, using a real database
// or the emulator as a backend, and with the option of multiple tests running side by side
// in multiple namespaces in a single collection to save costs.
//
// Configuration
//
// The standard configuration is to have a special file
// "testconfig.yaml" in the currenty directory when running the
// test. The config struct is expected inside a key "cosmostest", like this:
//
// cosmostest:
// Uri: "https://foo.documents.azure.com:443/"
// MasterKey: "yourkeyhere=="
// <... other fields from Config ...>
//
package cosmostest
import (
"context"
"crypto/tls"
"crypto/x509"
"fmt"
"net/http"
"github.com/adev73/go-cosmosdb/cosmos"
"github.com/adev73/go-cosmosdb/cosmosapi"
"github.com/adev73/go-cosmosdb/logging"
"github.com/gofrs/uuid"
"github.com/pkg/errors"
)
type Config struct {
Uri string `yaml:"Uri"`
MasterKey string `yaml:"MasterKey"`
MultiTenant bool `yaml:"MultiTenant"`
TlsCertificate string `yaml:"TlsCertificate"`
TlsServerName string `yaml:"TlsServerName"`
TlsInsecureSkipVerify bool `yaml:"TlsInsecureSkipVerify"`
DbName string `yaml:"DbName"`
CollectionIdPrefix string `yaml:"CollectionIdPrefix"`
AllowExistingCollection bool `yaml:"AllowExistingCollection"`
}
func check(err error, message string) {
if err != nil {
if message != "" {
err = errors.New(message)
}
panic(err)
}
}
// Factory for constructing the underlying, proper cosmosapi.Client given configuration.
// This is typically called by / wrapped by the test collection providers.
func RawClient(cfg Config) *cosmosapi.Client {
if cfg.Uri == "" {
panic("Missing requred parameter 'Uri'")
}
var caRoots *x509.CertPool
if cfg.TlsCertificate != "" {
caRoots = x509.NewCertPool()
if !caRoots.AppendCertsFromPEM([]byte(cfg.TlsCertificate)) {
panic("Failed to parse TLS certificate")
}
}
httpClient := &http.Client{
Transport: &http.Transport{
TLSClientConfig: &tls.Config{
RootCAs: caRoots,
ServerName: cfg.TlsServerName,
InsecureSkipVerify: cfg.TlsInsecureSkipVerify,
},
},
}
return cosmosapi.New(cfg.Uri, cosmosapi.Config{
MasterKey: cfg.MasterKey,
MaxRetries: 3,
}, httpClient, nil)
}
func SetupUniqueCollectionWithExistingDatabaseAndMinimalThroughput(log logging.StdLogger, cfg Config, id, partitionKey string) cosmos.Collection {
id = uuid.Must(uuid.NewV4()).String() + "-" + id
log.Printf("Creating Cosmos collection %s/%s\n", cfg.DbName, id)
client := RawClient(cfg)
_, err := client.CreateCollection(context.Background(), cfg.DbName, cosmosapi.CreateCollectionOptions{
Id: id,
PartitionKey: &cosmosapi.PartitionKey{
Paths: []string{"/" + partitionKey},
Kind: "Hash",
},
OfferThroughput: cosmosapi.OfferThroughput(400),
})
if err != nil {
panic(fmt.Sprintf("Failed to create Cosmos collection %s in database %s\n: %+v", id, cfg.DbName, err))
}
return cosmos.Collection{
Client: client,
DbName: cfg.DbName,
Name: id,
PartitionKey: partitionKey,
Context: context.Background(),
}
}
func SetupCollection(log logging.StdLogger, cfg Config, collectionId, partitionKey string) cosmos.Collection {
if cfg.CollectionIdPrefix == "" {
cfg.CollectionIdPrefix = uuid.Must(uuid.NewV4()).String() + "-"
}
if cfg.DbName == "" {
cfg.DbName = "default"
}
collectionId = cfg.CollectionIdPrefix + collectionId
client := RawClient(cfg)
if _, err := client.CreateDatabase(context.TODO(), cfg.DbName, nil); err != nil {
if errors.Cause(err) != cosmosapi.ErrConflict {
check(err, "Failed to create database")
}
// Database already existed, which is OK
}
_, err := client.CreateCollection(context.Background(), cfg.DbName, cosmosapi.CreateCollectionOptions{
Id: collectionId,
PartitionKey: &cosmosapi.PartitionKey{
Paths: []string{"/" + partitionKey},
Kind: "Hash",
},
OfferThroughput: 400,
})
if cfg.AllowExistingCollection && errors.Cause(err) == cosmosapi.ErrConflict {
err = nil
}
check(err, "")
return cosmos.Collection{
Client: client,
DbName: cfg.DbName,
Name: collectionId,
PartitionKey: partitionKey,
}
}
func TeardownCollection(collection cosmos.Collection) {
collection.Client.DeleteCollection(collection.Context, collection.DbName, collection.Name)
}