wallet/wallet_test.go

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