9 dbm "github.com/tendermint/tmlibs/db"
11 "github.com/bytom/account"
12 "github.com/bytom/asset"
13 "github.com/bytom/blockchain/pseudohsm"
14 "github.com/bytom/crypto/ed25519/chainkd"
15 "github.com/bytom/protocol"
16 "github.com/bytom/protocol/bc/types"
17 w "github.com/bytom/wallet"
20 type walletTestConfig struct {
21 Keys []*keyInfo `json:"keys"`
22 Accounts []*accountInfo `json:"accounts"`
23 Blocks []*wtBlock `json:"blocks"`
24 RollbackTo uint64 `json:"rollback_to"`
28 Name string `json:"name"`
29 Password string `json:"password"`
32 type accountInfo struct {
33 Name string `json:"name"`
34 Keys []string `json:"keys"`
35 Quorum int `json:"quorum"`
39 CoinbaseAccount string `json:"coinbase_account"`
40 Transactions []*wtTransaction `json:"transactions"`
41 PostStates []*accountBalance `json:"post_states"`
42 Append uint64 `json:"append"`
45 func (b *wtBlock) create(ctx *walletTestContext) (*types.Block, error) {
46 transactions := make([]*types.Tx, 0, len(b.Transactions))
47 for _, t := range b.Transactions {
48 tx, err := t.create(ctx)
52 transactions = append(transactions, tx)
54 return ctx.newBlock(transactions, b.CoinbaseAccount)
57 func (b *wtBlock) verifyPostStates(ctx *walletTestContext) error {
58 for _, state := range b.PostStates {
59 balance, err := ctx.getBalance(state.AccountAlias, state.AssetAlias)
64 if balance != state.Amount {
65 return fmt.Errorf("AccountAlias: %s, AssetAlias: %s, expected: %d, have: %d", state.AccountAlias, state.AssetAlias, state.Amount, balance)
71 type wtTransaction struct {
72 Passwords []string `json:"passwords"`
73 Inputs []*action `json:"inputs"`
74 Outputs []*action `json:"outputs"`
77 // create signed transaction
78 func (t *wtTransaction) create(ctx *walletTestContext) (*types.Tx, error) {
79 generator := NewTxGenerator(ctx.Wallet.AccountMgr, ctx.Wallet.AssetReg, ctx.Wallet.Hsm)
80 for _, input := range t.Inputs {
83 if err := generator.AddSpendInput(input.AccountAlias, input.AssetAlias, input.Amount); err != nil {
87 _, err := ctx.createAsset(input.AccountAlias, input.AssetAlias)
91 if err := generator.AddIssuanceInput(input.AssetAlias, input.Amount); err != nil {
97 for _, output := range t.Outputs {
100 if err := generator.AddTxOutput(output.AccountAlias, output.AssetAlias, output.Amount); err != nil {
104 if err := generator.AddRetirement(output.AssetAlias, output.Amount); err != nil {
109 return generator.Sign(t.Passwords)
113 Type string `json:"type"`
114 AccountAlias string `json:"name"`
115 AssetAlias string `json:"asset"`
116 Amount uint64 `json:"amount"`
119 type accountBalance struct {
120 AssetAlias string `json:"asset"`
121 AccountAlias string `json:"name"`
122 Amount uint64 `json:"amount"`
125 type walletTestContext struct {
127 Chain *protocol.Chain
130 func (ctx *walletTestContext) createControlProgram(accountName string, change bool) (*account.CtrlProgram, error) {
131 acc, err := ctx.Wallet.AccountMgr.FindByAlias(nil, accountName)
136 return ctx.Wallet.AccountMgr.CreateAddress(nil, acc.ID, change)
139 func (ctx *walletTestContext) getPubkey(keyAlias string) *chainkd.XPub {
140 pubKeys := ctx.Wallet.Hsm.ListKeys()
141 for i, key := range pubKeys {
142 if key.Alias == keyAlias {
143 return &pubKeys[i].XPub
149 func (ctx *walletTestContext) createAsset(accountAlias string, assetAlias string) (*asset.Asset, error) {
150 acc, err := ctx.Wallet.AccountMgr.FindByAlias(nil, accountAlias)
154 return ctx.Wallet.AssetReg.Define(acc.XPubs, len(acc.XPubs), nil, assetAlias)
157 func (ctx *walletTestContext) newBlock(txs []*types.Tx, coinbaseAccount string) (*types.Block, error) {
158 controlProgram, err := ctx.createControlProgram(coinbaseAccount, true)
162 return NewBlock(ctx.Chain, txs, controlProgram.ControlProgram)
165 func (ctx *walletTestContext) createKey(name string, password string) error {
166 _, err := ctx.Wallet.Hsm.XCreate(name, password)
170 func (ctx *walletTestContext) createAccount(name string, keys []string, quorum int) error {
171 xpubs := []chainkd.XPub{}
172 for _, alias := range keys {
173 xpub := ctx.getPubkey(alias)
175 return fmt.Errorf("can't find pubkey for %s", alias)
177 xpubs = append(xpubs, *xpub)
179 _, err := ctx.Wallet.AccountMgr.Create(nil, xpubs, quorum, name)
183 func (ctx *walletTestContext) update(block *types.Block) error {
184 if err := SolveAndUpdate(ctx.Chain, block); err != nil {
187 if err := ctx.Wallet.AttachBlock(block); err != nil {
193 func (ctx *walletTestContext) getBalance(accountAlias string, assetAlias string) (uint64, error) {
194 balances, _ := ctx.Wallet.GetAccountBalances("")
195 for _, balance := range balances {
196 if balance.Alias == accountAlias && balance.AssetAlias == assetAlias {
197 return balance.Amount, nil
203 func (ctx *walletTestContext) getAccBalances() map[string]map[string]uint64 {
204 accBalances := make(map[string]map[string]uint64)
205 balances, _ := ctx.Wallet.GetAccountBalances("")
206 for _, balance := range balances {
207 if accBalance, ok := accBalances[balance.Alias]; ok {
208 if _, ok := accBalance[balance.AssetAlias]; ok {
209 accBalance[balance.AssetAlias] += balance.Amount
212 accBalance[balance.AssetAlias] = balance.Amount
215 accBalances[balance.Alias] = map[string]uint64{balance.AssetAlias: balance.Amount}
220 func (ctx *walletTestContext) getDetachedBlocks(height uint64) ([]*types.Block, error) {
221 currentHeight := ctx.Chain.BestBlockHeight()
222 detachedBlocks := make([]*types.Block, 0, currentHeight-height)
223 for i := currentHeight; i > height; i-- {
224 block, err := ctx.Chain.GetBlockByHeight(i)
226 return detachedBlocks, err
228 detachedBlocks = append(detachedBlocks, block)
230 return detachedBlocks, nil
233 func (ctx *walletTestContext) validateRollback(oldAccBalances map[string]map[string]uint64) error {
234 accBalances := ctx.getAccBalances()
235 if reflect.DeepEqual(oldAccBalances, accBalances) {
238 return fmt.Errorf("different account balances after rollback")
241 func (cfg *walletTestConfig) Run() error {
242 dirPath, err := ioutil.TempDir(".", "pseudo_hsm")
246 defer os.RemoveAll(dirPath)
247 hsm, err := pseudohsm.New(dirPath)
252 db := dbm.NewDB("wallet_test_db", "leveldb", "wallet_test_db")
253 defer os.RemoveAll("wallet_test_db")
254 chain, _, _, _ := MockChain(db)
255 walletDB := dbm.NewDB("wallet", "leveldb", "wallet_db")
256 defer os.RemoveAll("wallet_db")
257 accountManager := account.NewManager(walletDB, chain)
258 assets := asset.NewRegistry(walletDB, chain)
259 wallet, err := w.NewWallet(walletDB, accountManager, assets, hsm, chain)
263 ctx := &walletTestContext{
268 for _, key := range cfg.Keys {
269 if err := ctx.createKey(key.Name, key.Password); err != nil {
274 for _, acc := range cfg.Accounts {
275 if err := ctx.createAccount(acc.Name, acc.Keys, acc.Quorum); err != nil {
280 var accBalances map[string]map[string]uint64
281 var rollbackBlock *types.Block
282 for _, blk := range cfg.Blocks {
283 block, err := blk.create(ctx)
287 if err := ctx.update(block); err != nil {
290 if err := blk.verifyPostStates(ctx); err != nil {
293 if block.Height <= cfg.RollbackTo && cfg.RollbackTo <= block.Height+blk.Append {
294 accBalances = ctx.getAccBalances()
295 rollbackBlock = block
297 if err := AppendBlocks(ctx.Chain, blk.Append); err != nil {
302 if rollbackBlock == nil {
306 // rollback and validate
307 detachedBlocks, err := ctx.getDetachedBlocks(rollbackBlock.Height)
312 forkedChain, err := declChain("forked_chain", ctx.Chain, rollbackBlock.Height, ctx.Chain.BestBlockHeight()+1)
313 defer os.RemoveAll("forked_chain")
318 if err := merge(forkedChain, ctx.Chain); err != nil {
322 for _, block := range detachedBlocks {
323 if err := ctx.Wallet.DetachBlock(block); err != nil {
327 return ctx.validateRollback(accBalances)