-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_builder.go
111 lines (95 loc) · 4.42 KB
/
test_builder.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
package app
import (
"github.com/cosmos/cosmos-sdk/codec"
sdk "github.com/cosmos/cosmos-sdk/types"
authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
vestingtypes "github.com/cosmos/cosmos-sdk/x/auth/vesting/types"
banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
)
// AuthBankGenesisBuilder is a tool for creating a combined auth and bank genesis state.
// Helper methods create basic accounts types and add them to a default genesis state.
// All methods return the builder so method calls can be chained together.
//
// Example:
//
// // create a single account genesis state
// builder := NewAuthBankGenesisBuilder().WithSimpleAccount(testUserAddress, testCoins)
// genesisState := builder.BuildMarshalled()
type AuthBankGenesisBuilder struct {
AuthGenesis authtypes.GenesisState
BankGenesis banktypes.GenesisState
}
// NewAuthBankGenesisBuilder creates a AuthBankGenesisBuilder containing default genesis states.
func NewAuthBankGenesisBuilder() *AuthBankGenesisBuilder {
return &AuthBankGenesisBuilder{
AuthGenesis: *authtypes.DefaultGenesisState(),
BankGenesis: *banktypes.DefaultGenesisState(),
}
}
// BuildMarshalled assembles the final GenesisState and json encodes it into a generic genesis type.
func (builder *AuthBankGenesisBuilder) BuildMarshalled(cdc codec.JSONCodec) GenesisState {
return GenesisState{
authtypes.ModuleName: cdc.MustMarshalJSON(&builder.AuthGenesis),
banktypes.ModuleName: cdc.MustMarshalJSON(&builder.BankGenesis),
}
}
// WithAccounts adds accounts of any type to the genesis state.
func (builder *AuthBankGenesisBuilder) WithAccounts(account ...authtypes.GenesisAccount) *AuthBankGenesisBuilder {
existing, err := authtypes.UnpackAccounts(builder.AuthGenesis.Accounts)
if err != nil {
panic(err)
}
existing = append(existing, account...)
existingPacked, err := authtypes.PackAccounts(existing)
if err != nil {
panic(err)
}
builder.AuthGenesis.Accounts = existingPacked
return builder
}
// WithBalances adds balances to the bank genesis state.
// It does not check the new denom is in the genesis state denom metadata.
func (builder *AuthBankGenesisBuilder) WithBalances(balance ...banktypes.Balance) *AuthBankGenesisBuilder {
builder.BankGenesis.Balances = append(builder.BankGenesis.Balances, balance...)
if !builder.BankGenesis.Supply.Empty() {
for _, b := range balance {
builder.BankGenesis.Supply = builder.BankGenesis.Supply.Add(b.Coins...)
}
}
return builder
}
// WithSimpleAccount adds a standard account to the genesis state.
func (builder *AuthBankGenesisBuilder) WithSimpleAccount(address sdk.AccAddress, balance sdk.Coins) *AuthBankGenesisBuilder {
return builder.
WithAccounts(authtypes.NewBaseAccount(address, nil, 0, 0)).
WithBalances(banktypes.Balance{Address: address.String(), Coins: balance})
}
// WithSimpleModuleAccount adds a module account to the genesis state.
func (builder *AuthBankGenesisBuilder) WithSimpleModuleAccount(moduleName string, balance sdk.Coins, permissions ...string) *AuthBankGenesisBuilder {
account := authtypes.NewEmptyModuleAccount(moduleName, permissions...)
return builder.
WithAccounts(account).
WithBalances(banktypes.Balance{Address: account.Address, Coins: balance})
}
// WithSimplePeriodicVestingAccount adds a periodic vesting account to the genesis state.
func (builder *AuthBankGenesisBuilder) WithSimplePeriodicVestingAccount(address sdk.AccAddress, balance sdk.Coins, periods vestingtypes.Periods, firstPeriodStartTimestamp int64) *AuthBankGenesisBuilder {
vestingAccount := newPeriodicVestingAccount(address, periods, firstPeriodStartTimestamp)
return builder.
WithAccounts(vestingAccount).
WithBalances(banktypes.Balance{Address: address.String(), Coins: balance})
}
// newPeriodicVestingAccount creates a periodic vesting account from a set of vesting periods.
func newPeriodicVestingAccount(address sdk.AccAddress, periods vestingtypes.Periods, firstPeriodStartTimestamp int64) *vestingtypes.PeriodicVestingAccount {
baseAccount := authtypes.NewBaseAccount(address, nil, 0, 0)
originalVesting := sdk.NewCoins()
for _, p := range periods {
originalVesting = originalVesting.Add(p.Amount...)
}
var totalPeriods int64
for _, p := range periods {
totalPeriods += p.Length
}
endTime := firstPeriodStartTimestamp + totalPeriods
baseVestingAccount := vestingtypes.NewBaseVestingAccount(baseAccount, originalVesting, endTime)
return vestingtypes.NewPeriodicVestingAccountRaw(baseVestingAccount, firstPeriodStartTimestamp, periods)
}