test/wallet_test.go

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