wallet/wallet_test.go

   1 package wallet
   2
   3 import (
   4         "io/ioutil"
   5         "os"
   6         "testing"
   7         "time"
   8
   9         dbm "github.com/tendermint/tmlibs/db"
  10
  11         "github.com/bytom/account"
  12         "github.com/bytom/asset"
  13         "github.com/bytom/blockchain/pseudohsm"
  14         "github.com/bytom/blockchain/signers"
  15         "github.com/bytom/blockchain/txbuilder"
  16         "github.com/bytom/consensus"
  17         "github.com/bytom/crypto/ed25519/chainkd"
  18         "github.com/bytom/database/leveldb"
  19         "github.com/bytom/event"
  20         "github.com/bytom/protocol"
  21         "github.com/bytom/protocol/bc"
  22         "github.com/bytom/protocol/bc/types"
  23 )
  24
  25 func TestWalletUpdate(t *testing.T) {
  26         dirPath, err := ioutil.TempDir(".", "")
  27         if err != nil {
  28                 t.Fatal(err)
  29         }
  30         defer os.RemoveAll(dirPath)
  31
  32         testDB := dbm.NewDB("testdb", "leveldb", "temp")
  33         defer os.RemoveAll("temp")
  34
  35         store := leveldb.NewStore(testDB)
  36         dispatcher := event.NewDispatcher()
  37         txPool := protocol.NewTxPool(store, dispatcher)
  38
  39         chain, err := protocol.NewChain(store, txPool)
  40         if err != nil {
  41                 t.Fatal(err)
  42         }
  43
  44         accountManager := account.NewManager(testDB, chain)
  45         hsm, err := pseudohsm.New(dirPath)
  46         if err != nil {
  47                 t.Fatal(err)
  48         }
  49
  50         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
  51         if err != nil {
  52                 t.Fatal(err)
  53         }
  54
  55         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
  56         if err != nil {
  57                 t.Fatal(err)
  58         }
  59
  60         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
  61         if err != nil {
  62                 t.Fatal(err)
  63         }
  64
  65         controlProg.KeyIndex = 1
  66
  67         reg := asset.NewRegistry(testDB, chain)
  68         asset, err := reg.Define([]chainkd.XPub{xpub1.XPub}, 1, nil, "TESTASSET", nil)
  69         if err != nil {
  70                 t.Fatal(err)
  71         }
  72
  73         utxos := []*account.UTXO{}
  74         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
  75         utxos = append(utxos, btmUtxo)
  76         OtherUtxo := mockUTXO(controlProg, &asset.AssetID)
  77         utxos = append(utxos, OtherUtxo)
  78
  79         _, txData, err := mockTxData(utxos, testAccount)
  80         if err != nil {
  81                 t.Fatal(err)
  82         }
  83
  84         tx := types.NewTx(*txData)
  85         block := mockSingleBlock(tx)
  86         txStatus := bc.NewTransactionStatus()
  87         txStatus.SetStatus(0, false)
  88         store.SaveBlock(block, txStatus)
  89
  90         w := mockWallet(testDB, accountManager, reg, chain, dispatcher)
  91         err = w.AttachBlock(block)
  92         if err != nil {
  93                 t.Fatal(err)
  94         }
  95
  96         if _, err := w.GetTransactionByTxID(tx.ID.String()); err != nil {
  97                 t.Fatal(err)
  98         }
  99
 100         wants, err := w.GetTransactions("")
 101         if len(wants) != 1 {
 102                 t.Fatal(err)
 103         }
 104 }
 105
 106 func TestMemPoolTxQueryLoop(t *testing.T) {
 107         dirPath, err := ioutil.TempDir(".", "")
 108         if err != nil {
 109                 t.Fatal(err)
 110         }
 111         defer os.RemoveAll(dirPath)
 112
 113         testDB := dbm.NewDB("testdb", "leveldb", dirPath)
 114
 115         store := leveldb.NewStore(testDB)
 116         dispatcher := event.NewDispatcher()
 117         txPool := protocol.NewTxPool(store, dispatcher)
 118
 119         chain, err := protocol.NewChain(store, txPool)
 120         if err != nil {
 121                 t.Fatal(err)
 122         }
 123
 124         accountManager := account.NewManager(testDB, chain)
 125         hsm, err := pseudohsm.New(dirPath)
 126         if err != nil {
 127                 t.Fatal(err)
 128         }
 129
 130         xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
 131         if err != nil {
 132                 t.Fatal(err)
 133         }
 134
 135         testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
 136         if err != nil {
 137                 t.Fatal(err)
 138         }
 139
 140         controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
 141         if err != nil {
 142                 t.Fatal(err)
 143         }
 144
 145         controlProg.KeyIndex = 1
 146
 147         reg := asset.NewRegistry(testDB, chain)
 148         asset, err := reg.Define([]chainkd.XPub{xpub1.XPub}, 1, nil, "TESTASSET", nil)
 149         if err != nil {
 150                 t.Fatal(err)
 151         }
 152
 153         utxos := []*account.UTXO{}
 154         btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
 155         utxos = append(utxos, btmUtxo)
 156         OtherUtxo := mockUTXO(controlProg, &asset.AssetID)
 157         utxos = append(utxos, OtherUtxo)
 158
 159         _, txData, err := mockTxData(utxos, testAccount)
 160         if err != nil {
 161                 t.Fatal(err)
 162         }
 163
 164         tx := types.NewTx(*txData)
 165         //block := mockSingleBlock(tx)
 166         txStatus := bc.NewTransactionStatus()
 167         txStatus.SetStatus(0, false)
 168         w, err := NewWallet(testDB, accountManager, reg, hsm, chain, dispatcher)
 169         go w.memPoolTxQueryLoop()
 170         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgNewTx}})
 171         time.Sleep(time.Millisecond * 10)
 172         if _, err = w.GetUnconfirmedTxByTxID(tx.ID.String()); err != nil {
 173                 t.Fatal("disaptch new tx msg error:", err)
 174         }
 175         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: protocol.MsgRemoveTx}})
 176         time.Sleep(time.Millisecond * 10)
 177         txs, err := w.GetUnconfirmedTxs(testAccount.ID)
 178         if err != nil {
 179                 t.Fatal("get unconfirmed tx error:", err)
 180         }
 181
 182         if len(txs) != 0 {
 183                 t.Fatal("disaptch remove tx msg error")
 184         }
 185
 186         w.eventDispatcher.Post(protocol.TxMsgEvent{TxMsg: &protocol.TxPoolMsg{TxDesc: &protocol.TxDesc{Tx: tx}, MsgType: 2}})
 187 }
 188
 189 func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
 190         utxo := &account.UTXO{}
 191         utxo.OutputID = bc.Hash{V0: 1}
 192         utxo.SourceID = bc.Hash{V0: 2}
 193         utxo.AssetID = *assetID
 194         utxo.Amount = 1000000000
 195         utxo.SourcePos = 0
 196         utxo.ControlProgram = controlProg.ControlProgram
 197         utxo.AccountID = controlProg.AccountID
 198         utxo.Address = controlProg.Address
 199         utxo.ControlProgramIndex = controlProg.KeyIndex
 200         return utxo
 201 }
 202
 203 func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
 204         tplBuilder := txbuilder.NewBuilder(time.Now())
 205
 206         for _, utxo := range utxos {
 207                 txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
 208                 if err != nil {
 209                         return nil, nil, err
 210                 }
 211                 tplBuilder.AddInput(txInput, sigInst)
 212
 213                 out := &types.TxOutput{}
 214                 if utxo.AssetID == *consensus.BTMAssetID {
 215                         out = types.NewTxOutput(utxo.AssetID, 100, utxo.ControlProgram)
 216                 } else {
 217                         out = types.NewTxOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
 218                 }
 219                 tplBuilder.AddOutput(out)
 220         }
 221
 222         return tplBuilder.Build()
 223 }
 224
 225 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain, dispatcher *event.Dispatcher) *Wallet {
 226         wallet := &Wallet{
 227                 DB:              walletDB,
 228                 AccountMgr:      account,
 229                 AssetReg:        asset,
 230                 chain:           chain,
 231                 RecoveryMgr:     newRecoveryManager(walletDB, account),
 232                 eventDispatcher: dispatcher,
 233         }
 234         wallet.txMsgSub, _ = wallet.eventDispatcher.Subscribe(protocol.TxMsgEvent{})
 235         return wallet
 236 }
 237
 238 func mockSingleBlock(tx *types.Tx) *types.Block {
 239         return &types.Block{
 240                 BlockHeader: types.BlockHeader{
 241                         Version: 1,
 242                         Height:  1,
 243                         Bits:    2305843009230471167,
 244                 },
 245                 Transactions: []*types.Tx{tx},
 246         }
 247 }