test/wallet_test.go

   1 package test
   2
   3 import (
   4         "io/ioutil"
   5         "os"
   6         "reflect"
   7         "testing"
   8         "time"
   9
  10         "github.com/bytom/vapor/account"
  11         "github.com/bytom/vapor/asset"
  12         "github.com/bytom/vapor/blockchain/pseudohsm"
  13         "github.com/bytom/vapor/blockchain/signers"
  14         "github.com/bytom/vapor/blockchain/txbuilder"
  15         "github.com/bytom/vapor/config"
  16         "github.com/bytom/vapor/consensus"
  17         "github.com/bytom/vapor/crypto/ed25519/chainkd"
  18         "github.com/bytom/vapor/database"
  19         dbm "github.com/bytom/vapor/database/leveldb"
  20         "github.com/bytom/vapor/event"
  21         "github.com/bytom/vapor/protocol"
  22         "github.com/bytom/vapor/protocol/bc"
  23         "github.com/bytom/vapor/protocol/bc/types"
  24         wt "github.com/bytom/vapor/wallet"
  25 )
  26
  27 func TestWalletUpdate(t *testing.T) {
  28         dirPath, err := ioutil.TempDir(".", "")
  29         if err != nil {
  30                 t.Fatal(err)
  31         }
  32         defer os.RemoveAll(dirPath)
  33
  34         config.CommonConfig = config.DefaultConfig()
  35         testDB := dbm.NewDB("testdb", "leveldb", "temp")
  36         defer func() {
  37                 testDB.Close()
  38                 os.RemoveAll("temp")
  39         }()
  40
  41         store := database.NewStore(testDB)
  42         walletStore := database.NewWalletStore(testDB)
  43         dispatcher := event.NewDispatcher()
  44         txPool := protocol.NewTxPool(store, dispatcher)
  45
  46         chain, err := protocol.NewChain(store, txPool, dispatcher)
  47         if err != nil {
  48                 t.Fatal(err)
  49         }
  50
  51         accountStore := database.NewAccountStore(testDB)
  52         accountManager := account.NewManager(accountStore, chain)
  53         hsm, err := pseudohsm.New(dirPath)
  54         if err != nil {
  55                 t.Fatal(err)
  56         }
  57
  58         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
  59         if err != nil {
  60                 t.Fatal(err)
  61         }
  62
  63         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
  64         if err != nil {
  65                 t.Fatal(err)
  66         }
  67
  68         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
  69         if err != nil {
  70                 t.Fatal(err)
  71         }
  72
  73         controlProg.KeyIndex = 1
  74
  75         reg := asset.NewRegistry(testDB, chain)
  76         asset := bc.AssetID{V0: 5}
  77
  78         utxos := []*account.UTXO{}
  79         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
  80         utxos = append(utxos, btmUtxo)
  81         OtherUtxo := mockUTXO(controlProg, &asset)
  82         utxos = append(utxos, OtherUtxo)
  83
  84         _, txData, err := mockTxData(utxos, testAccount)
  85         if err != nil {
  86                 t.Fatal(err)
  87         }
  88
  89         tx := types.NewTx(*txData)
  90         block := mockSingleBlock(tx)
  91         txStatus := bc.NewTransactionStatus()
  92         txStatus.SetStatus(0, false)
  93         txStatus.SetStatus(1, false)
  94         store.SaveBlock(block, txStatus)
  95
  96         w := newMockWallet(walletStore, accountManager, reg, chain, dispatcher, true)
  97         err = w.AttachBlock(block)
  98         if err != nil {
  99                 t.Fatal(err)
 100         }
 101
 102         if _, err := w.GetTransactionByTxID(tx.ID.String()); err != nil {
 103                 t.Fatal(err)
 104         }
 105
 106         wants, err := w.GetTransactions(testAccount.ID, "", 1, false)
 107         if len(wants) != 1 {
 108                 t.Fatal(err)
 109         }
 110
 111         if wants[0].ID != tx.ID {
 112                 t.Fatal("account txID mismatch")
 113         }
 114
 115         for position, tx := range block.Transactions {
 116                 get := testDB.Get(database.CalcGlobalTxIndexKey(tx.ID.String()))
 117                 bh := block.BlockHeader.Hash()
 118                 expect := database.CalcGlobalTxIndex(&bh, uint64(position))
 119                 if !reflect.DeepEqual(get, expect) {
 120                         t.Fatalf("position#%d: compare retrieved globalTxIdx err", position)
 121                 }
 122         }
 123 }
 124
 125 func TestRescanWallet(t *testing.T) {
 126         // prepare wallet & db.
 127         dirPath, err := ioutil.TempDir(".", "")
 128         if err != nil {
 129                 t.Fatal(err)
 130         }
 131         defer os.RemoveAll(dirPath)
 132
 133         config.CommonConfig = config.DefaultConfig()
 134         testDB := dbm.NewDB("testdb", "leveldb", "temp")
 135         walletStore := database.NewWalletStore(testDB)
 136         defer func() {
 137                 testDB.Close()
 138                 os.RemoveAll("temp")
 139         }()
 140
 141         store := database.NewStore(testDB)
 142         dispatcher := event.NewDispatcher()
 143         txPool := protocol.NewTxPool(store, dispatcher)
 144         chain, err := protocol.NewChain(store, txPool, dispatcher)
 145         if err != nil {
 146                 t.Fatal(err)
 147         }
 148
 149         statusInfo := wt.StatusInfo{
 150                 Version:  uint(1),
 151                 WorkHash: bc.Hash{V0: 0xff},
 152         }
 153         if err := walletStore.SetWalletInfo(&statusInfo); err != nil {
 154                 t.Fatal(err)
 155         }
 156         walletInfo, err := walletStore.GetWalletInfo()
 157         if err != nil {
 158                 t.Fatal(err)
 159         }
 160
 161         accountStore := database.NewAccountStore(testDB)
 162         accountManager := account.NewManager(accountStore, chain)
 163         w := newMockWallet(walletStore, accountManager, nil, chain, dispatcher, false)
 164         if err != nil {
 165                 t.Fatal(err)
 166         }
 167         w.Status = *walletInfo
 168
 169         // rescan wallet.
 170         if err := w.LoadWalletInfo(); err != nil {
 171                 t.Fatal(err)
 172         }
 173
 174         block := config.GenesisBlock()
 175         if w.Status.WorkHash != block.Hash() {
 176                 t.Fatal("reattach from genesis block")
 177         }
 178 }
 179
 180 func TestMemPoolTxQueryLoop(t *testing.T) {
 181         dirPath, err := ioutil.TempDir(".", "")
 182         if err != nil {
 183                 t.Fatal(err)
 184         }
 185         config.CommonConfig = config.DefaultConfig()
 186         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 187         defer func() {
 188                 testDB.Close()
 189                 os.RemoveAll(dirPath)
 190         }()
 191
 192         store := database.NewStore(testDB)
 193         dispatcher := event.NewDispatcher()
 194         txPool := protocol.NewTxPool(store, dispatcher)
 195
 196         chain, err := protocol.NewChain(store, txPool, dispatcher)
 197         if err != nil {
 198                 t.Fatal(err)
 199         }
 200
 201         accountStore := database.NewAccountStore(testDB)
 202         accountManager := account.NewManager(accountStore, chain)
 203         hsm, err := pseudohsm.New(dirPath)
 204         if err != nil {
 205                 t.Fatal(err)
 206         }
 207
 208         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 209         if err != nil {
 210                 t.Fatal(err)
 211         }
 212
 213         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
 214         if err != nil {
 215                 t.Fatal(err)
 216         }
 217
 218         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
 219         if err != nil {
 220                 t.Fatal(err)
 221         }
 222
 223         controlProg.KeyIndex = 1
 224
 225         reg := asset.NewRegistry(testDB, chain)
 226         asset := bc.AssetID{V0: 5}
 227
 228         utxos := []*account.UTXO{}
 229         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
 230         utxos = append(utxos, btmUtxo)
 231         OtherUtxo := mockUTXO(controlProg, &asset)
 232         utxos = append(utxos, OtherUtxo)
 233
 234         _, txData, err := mockTxData(utxos, testAccount)
 235         if err != nil {
 236                 t.Fatal(err)
 237         }
 238
 239         tx := types.NewTx(*txData)
 240         txStatus := bc.NewTransactionStatus()
 241         txStatus.SetStatus(0, false)
 242         walletStore := database.NewWalletStore(testDB)
 243         w := newMockWallet(walletStore, accountManager, reg, chain, dispatcher, false)
 244         go w.MemPoolTxQueryLoop()
 245         w.EventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgNewTx}})
 246         time.Sleep(time.Millisecond * 10)
 247         if _, err := w.GetUnconfirmedTxByTxID(tx.ID.String()); err != nil {
 248                 t.Fatal("dispatch new tx msg error:", err)
 249         }
 250         w.EventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgRemoveTx}})
 251         time.Sleep(time.Millisecond * 10)
 252         txs, err := w.GetUnconfirmedTxs(testAccount.ID)
 253         if err != nil {
 254                 t.Fatal("get unconfirmed tx error:", err)
 255         }
 256
 257         if len(txs) != 0 {
 258                 t.Fatal("dispatch remove tx msg error")
 259         }
 260
 261         w.EventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: 2}})
 262 }
 263
 264 func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
 265         utxo := &account.UTXO{}
 266         utxo.OutputID = bc.Hash{V0: 1}
 267         utxo.SourceID = bc.Hash{V0: 2}
 268         utxo.AssetID = *assetID
 269         utxo.Amount = 1000000000
 270         utxo.SourcePos = 0
 271         utxo.ControlProgram = controlProg.ControlProgram
 272         utxo.AccountID = controlProg.AccountID
 273         utxo.Address = controlProg.Address
 274         utxo.ControlProgramIndex = controlProg.KeyIndex
 275         return utxo
 276 }
 277
 278 func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
 279         tplBuilder := txbuilder.NewBuilder(time.Now())
 280
 281         for _, utxo := range utxos {
 282                 txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
 283                 if err != nil {
 284                         return nil, nil, err
 285                 }
 286                 tplBuilder.AddInput(txInput, sigInst)
 287
 288                 out := &types.TxOutput{}
 289                 if utxo.AssetID == *consensus.BTMAssetID {
 290                         out = types.NewIntraChainOutput(utxo.AssetID, 100, utxo.ControlProgram)
 291                 } else {
 292                         out = types.NewIntraChainOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
 293                 }
 294                 tplBuilder.AddOutput(out)
 295         }
 296
 297         return tplBuilder.Build()
 298 }
 299
 300 type mockWallet struct {
 301         Wallet *wt.Wallet
 302 }
 303
 304 func newMockWallet(store wt.WalletStore, account *account.Manager, asset *asset.Registry, chain *protocol.Chain, dispatcher *event.Dispatcher, txIndexFlag bool) *wt.Wallet {
 305         wallet := &wt.Wallet{
 306                 Store:           store,
 307                 AccountMgr:      account,
 308                 AssetReg:        asset,
 309                 Chain:           chain,
 310                 RecoveryMgr:     wt.NewRecoveryManager(store, account),
 311                 EventDispatcher: dispatcher,
 312                 TxIndexFlag:     txIndexFlag,
 313         }
 314         wallet.TxMsgSub, _ = wallet.EventDispatcher.Subscribe(protocol.TxMsgEvent{})
 315         return wallet
 316 }
 317
 318 func mockSingleBlock(tx *types.Tx) *types.Block {
 319         return &types.Block{
 320                 BlockHeader: types.BlockHeader{
 321                         Version: 1,
 322                         Height:  1,
 323                 },
 324                 Transactions: []*types.Tx{config.GenesisTx(), tx},
 325         }
 326 }