-
Notifications
You must be signed in to change notification settings - Fork 664
/
env.go
208 lines (181 loc) · 6.66 KB
/
env.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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
// Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package e2e
import (
"encoding/json"
"errors"
"math/rand"
"os"
"time"
"github.com/stretchr/testify/require"
"github.com/ava-labs/avalanchego/api/info"
"github.com/ava-labs/avalanchego/config"
"github.com/ava-labs/avalanchego/tests"
"github.com/ava-labs/avalanchego/tests/fixture"
"github.com/ava-labs/avalanchego/tests/fixture/tmpnet"
"github.com/ava-labs/avalanchego/utils/crypto/secp256k1"
"github.com/ava-labs/avalanchego/vms/secp256k1fx"
ginkgo "github.com/onsi/ginkgo/v2"
)
// Env is used to access shared test fixture. Intended to be
// initialized from SynchronizedBeforeSuite.
var Env *TestEnvironment
func InitSharedTestEnvironment(envBytes []byte) {
require := require.New(ginkgo.GinkgoT())
require.Nil(Env, "env already initialized")
Env = &TestEnvironment{}
require.NoError(json.Unmarshal(envBytes, Env))
Env.require = require
}
type TestEnvironment struct {
// The directory where the test network configuration is stored
NetworkDir string
// URIs used to access the API endpoints of nodes of the network
URIs []tmpnet.NodeURI
// The URI used to access the http server that allocates test data
TestDataServerURI string
// The duration to wait before shutting down private networks. A
// non-zero value may be useful to ensure all metrics can be
// scraped before shutdown.
PrivateNetworkShutdownDelay time.Duration
require *require.Assertions
}
func (te *TestEnvironment) Marshal() []byte {
bytes, err := json.Marshal(te)
require.NoError(ginkgo.GinkgoT(), err)
return bytes
}
// Initialize a new test environment with a shared network (either pre-existing or newly created).
func NewTestEnvironment(flagVars *FlagVars, desiredNetwork *tmpnet.Network) *TestEnvironment {
require := require.New(ginkgo.GinkgoT())
var network *tmpnet.Network
// Need to load the network if it is being stopped or reused
if flagVars.StopNetwork() || flagVars.ReuseNetwork() {
networkDir := flagVars.NetworkDir()
var networkSymlink string // If populated, prompts removal of the referenced symlink if --stop-network is specified
if len(networkDir) == 0 {
// Attempt to reuse the network at the default owner path
symlinkPath, err := tmpnet.GetReusableNetworkPathForOwner(desiredNetwork.Owner)
require.NoError(err)
_, err = os.Stat(symlinkPath)
if !errors.Is(err, os.ErrNotExist) {
// Try to load the existing network
require.NoError(err)
networkDir = symlinkPath
// Enable removal of the referenced symlink if --stop-network is specified
networkSymlink = symlinkPath
}
}
if len(networkDir) > 0 {
var err error
network, err = tmpnet.ReadNetwork(networkDir)
require.NoError(err)
tests.Outf("{{yellow}}Loaded a network configured at %s{{/}}\n", network.Dir)
}
if flagVars.StopNetwork() {
if len(networkSymlink) > 0 {
// Remove the symlink to avoid attempts to reuse the stopped network
tests.Outf("Removing symlink %s\n", networkSymlink)
if err := os.Remove(networkSymlink); !errors.Is(err, os.ErrNotExist) {
require.NoError(err)
}
}
if network != nil {
tests.Outf("Stopping network\n")
require.NoError(network.Stop(DefaultContext()))
} else {
tests.Outf("No network to stop\n")
}
os.Exit(0)
}
}
// Start a new network
if network == nil {
network = desiredNetwork
StartNetwork(
network,
flagVars.AvalancheGoExecPath(),
flagVars.PluginDir(),
flagVars.NetworkShutdownDelay(),
flagVars.ReuseNetwork(),
)
// Wait for chains to have bootstrapped on all nodes
Eventually(func() bool {
for _, subnet := range network.Subnets {
for _, validatorID := range subnet.ValidatorIDs {
uri, err := network.GetURIForNodeID(validatorID)
require.NoError(err)
infoClient := info.NewClient(uri)
for _, chain := range subnet.Chains {
isBootstrapped, err := infoClient.IsBootstrapped(DefaultContext(), chain.ChainID.String())
// Ignore errors since a chain id that is not yet known will result in a recoverable error.
if err != nil || !isBootstrapped {
return false
}
}
}
}
return true
}, DefaultTimeout, DefaultPollingInterval, "failed to see all chains bootstrap before timeout")
}
uris := network.GetNodeURIs()
require.NotEmpty(uris, "network contains no nodes")
tests.Outf("{{green}}network URIs: {{/}} %+v\n", uris)
testDataServerURI, err := fixture.ServeTestData(fixture.TestData{
PreFundedKeys: network.PreFundedKeys,
})
tests.Outf("{{green}}test data server URI: {{/}} %+v\n", testDataServerURI)
require.NoError(err)
return &TestEnvironment{
NetworkDir: network.Dir,
URIs: uris,
TestDataServerURI: testDataServerURI,
PrivateNetworkShutdownDelay: flagVars.NetworkShutdownDelay(),
require: require,
}
}
// Retrieve a random URI to naively attempt to spread API load across
// nodes.
func (te *TestEnvironment) GetRandomNodeURI() tmpnet.NodeURI {
r := rand.New(rand.NewSource(time.Now().Unix())) //#nosec G404
nodeURI := te.URIs[r.Intn(len(te.URIs))]
tests.Outf("{{blue}} targeting node %s with URI: %s{{/}}\n", nodeURI.NodeID, nodeURI.URI)
return nodeURI
}
// Retrieve the network to target for testing.
func (te *TestEnvironment) GetNetwork() *tmpnet.Network {
network, err := tmpnet.ReadNetwork(te.NetworkDir)
te.require.NoError(err)
return network
}
// Retrieve the specified number of funded keys allocated for the caller's exclusive use.
func (te *TestEnvironment) AllocatePreFundedKeys(count int) []*secp256k1.PrivateKey {
keys, err := fixture.AllocatePreFundedKeys(te.TestDataServerURI, count)
te.require.NoError(err)
tests.Outf("{{blue}} allocated pre-funded key(s): %+v{{/}}\n", keys)
return keys
}
// Retrieve a funded key allocated for the caller's exclusive use.
func (te *TestEnvironment) AllocatePreFundedKey() *secp256k1.PrivateKey {
return te.AllocatePreFundedKeys(1)[0]
}
// Create a new keychain with the specified number of test keys.
func (te *TestEnvironment) NewKeychain(count int) *secp256k1fx.Keychain {
keys := te.AllocatePreFundedKeys(count)
return secp256k1fx.NewKeychain(keys...)
}
// Create a new private network that is not shared with other tests.
func (te *TestEnvironment) StartPrivateNetwork(network *tmpnet.Network) {
// Use the same configuration as the shared network
sharedNetwork, err := tmpnet.ReadNetwork(te.NetworkDir)
te.require.NoError(err)
pluginDir, err := sharedNetwork.DefaultFlags.GetStringVal(config.PluginDirKey)
te.require.NoError(err)
StartNetwork(
network,
sharedNetwork.DefaultRuntimeConfig.AvalancheGoPath,
pluginDir,
te.PrivateNetworkShutdownDelay,
false, /* reuseNetwork */
)
}