OSDN Git Service

fix(api): delete transactions and rescan blocks after updating account alias (#1466)
[bytom/bytom.git] / wallet / wallet_test.go
index 13bf241..1914016 100644 (file)
@@ -1,22 +1,20 @@
 package wallet
 
 import (
-       "context"
        "io/ioutil"
        "os"
        "testing"
        "time"
 
-       "github.com/tendermint/go-wire/data/base58"
        dbm "github.com/tendermint/tmlibs/db"
 
        "github.com/bytom/account"
        "github.com/bytom/asset"
        "github.com/bytom/blockchain/pseudohsm"
+       "github.com/bytom/blockchain/signers"
        "github.com/bytom/blockchain/txbuilder"
        "github.com/bytom/consensus"
        "github.com/bytom/crypto/ed25519/chainkd"
-       "github.com/bytom/crypto/sha3pool"
        "github.com/bytom/database/leveldb"
        "github.com/bytom/protocol"
        "github.com/bytom/protocol/bc"
@@ -34,7 +32,7 @@ func TestWalletUpdate(t *testing.T) {
        defer os.RemoveAll("temp")
 
        store := leveldb.NewStore(testDB)
-       txPool := protocol.NewTxPool()
+       txPool := protocol.NewTxPool(store)
 
        chain, err := protocol.NewChain(store, txPool)
        if err != nil {
@@ -47,160 +45,67 @@ func TestWalletUpdate(t *testing.T) {
                t.Fatal(err)
        }
 
-       xpub1, err := hsm.XCreate("test_pub1", "password")
+       xpub1, _, err := hsm.XCreate("test_pub1", "password", "en")
        if err != nil {
                t.Fatal(err)
        }
 
-       testAccount, err := accountManager.Create(nil, []chainkd.XPub{xpub1.XPub}, 1, "testAccount", nil)
+       testAccount, err := accountManager.Create([]chainkd.XPub{xpub1.XPub}, 1, "testAccount", signers.BIP0044)
        if err != nil {
                t.Fatal(err)
        }
 
-       controlProg, err := accountManager.CreateAddress(nil, testAccount.ID, false)
+       controlProg, err := accountManager.CreateAddress(testAccount.ID, false)
        if err != nil {
                t.Fatal(err)
        }
 
        controlProg.KeyIndex = 1
 
-       utxo := mockUTXO(controlProg)
-       _, txData, err := mockTxData(utxo, testAccount)
-       if err != nil {
-               t.Fatal(err)
-       }
-
-       tx := types.NewTx(*txData)
-
        reg := asset.NewRegistry(testDB, chain)
-
-       w := mockWallet(testDB, accountManager, reg, chain)
-
-       block := mockSingleBlock(tx)
-
-       txStatus := bc.NewTransactionStatus()
-       store.SaveBlock(block, txStatus)
-
-       err = w.attachBlock(block)
+       asset, err := reg.Define([]chainkd.XPub{xpub1.XPub}, 1, nil, "TESTASSET", nil)
        if err != nil {
                t.Fatal(err)
        }
 
-       want, err := w.GetTransactionsByTxID(tx.ID.String())
-       if len(want) != 1 {
-               t.Fatal(err)
-       }
-
-       wants, err := w.GetTransactionsByTxID("")
-       if len(wants) != 1 {
-               t.Fatal(err)
-       }
-}
+       utxos := []*account.UTXO{}
+       btmUtxo := mockUTXO(controlProg, consensus.BTMAssetID)
+       utxos = append(utxos, btmUtxo)
+       OtherUtxo := mockUTXO(controlProg, &asset.AssetID)
+       utxos = append(utxos, OtherUtxo)
 
-func TestExportAndImportPrivKey(t *testing.T) {
-       dirPath, err := ioutil.TempDir(".", "")
-       if err != nil {
-               t.Fatal(err)
-       }
-       defer os.RemoveAll(dirPath)
-
-       testDB := dbm.NewDB("testdb", "leveldb", "temp")
-       defer os.RemoveAll("temp")
-
-       store := leveldb.NewStore(testDB)
-       txPool := protocol.NewTxPool()
-
-       chain, err := protocol.NewChain(store, txPool)
-       if err != nil {
-               t.Fatal(err)
-       }
-
-       acntManager := account.NewManager(testDB, chain)
-       reg := asset.NewRegistry(testDB, chain)
-
-       hsm, err := pseudohsm.New(dirPath)
+       _, txData, err := mockTxData(utxos, testAccount)
        if err != nil {
                t.Fatal(err)
        }
 
-       pwd := "password"
-       xpub, err := hsm.XCreate("alias", pwd)
-       if err != nil {
-               t.Fatal(err)
-       }
-
-       w, err := NewWallet(testDB, acntManager, reg, hsm, chain)
-       if err != nil {
-               t.Fatal(err)
-       }
-
-       ctx := context.Background()
-       acnt1, err := w.AccountMgr.Create(ctx, []chainkd.XPub{xpub.XPub}, 1, "account-alias", nil)
-       if err != nil {
-               t.Fatal(err)
-       }
-
-       priv, err := w.ExportAccountPrivKey(xpub.XPub, pwd)
-
-       wantPriv, err := hsm.LoadChainKDKey(xpub.XPub, pwd)
-       if err != nil {
-               t.Fatal(err)
-       }
-       var hashed [32]byte
-       sha3pool.Sum256(hashed[:], wantPriv[:])
-
-       tmp := append(wantPriv[:], hashed[:4]...)
-       res := base58.Encode(tmp)
-
-       if res != *priv {
-               t.Fatalf("XPrivs should be identical.\nBefore: %v\n After: %v\n", *priv, res)
-       }
+       tx := types.NewTx(*txData)
+       block := mockSingleBlock(tx)
+       txStatus := bc.NewTransactionStatus()
+       txStatus.SetStatus(0, false)
+       store.SaveBlock(block, txStatus)
 
-       rawPriv, err := base58.Decode(*priv)
+       w := mockWallet(testDB, accountManager, reg, chain)
+       err = w.AttachBlock(block)
        if err != nil {
                t.Fatal(err)
        }
 
-       if len(rawPriv) != 68 {
-               t.Fatal("invalid private key hash length")
-       }
-
-       var xprv [64]byte
-       copy(xprv[:], rawPriv[:64])
-
-       _, err = w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
-       if err != pseudohsm.ErrDuplicateKeyAlias {
-               t.Fatal(err)
-       }
-
-       hsm.XDelete(xpub.XPub, pwd)
-
-       _, err = w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
-       if err != account.ErrDuplicateAlias {
+       if _, err := w.GetTransactionByTxID(tx.ID.String()); err != nil {
                t.Fatal(err)
        }
 
-       accountInfo := struct {
-               AccountInfo string `json:"account_info"`
-       }{AccountInfo: acnt1.Alias}
-
-       w.AccountMgr.DeleteAccount(accountInfo)
-
-       acnt2, err := w.ImportAccountPrivKey(xprv, xpub.Alias, pwd, 0, acnt1.Alias)
-       if err != nil {
+       wants, err := w.GetTransactions("")
+       if len(wants) != 1 {
                t.Fatal(err)
        }
-
-       if acnt2.XPub != acnt1.XPubs[0] {
-               t.Fatalf("XPubs should be identical.\nBefore: %v\n After: %v\n", acnt1.XPubs[0], acnt2.XPub)
-       }
 }
 
-func mockUTXO(controlProg *account.CtrlProgram) *account.UTXO {
+func mockUTXO(controlProg *account.CtrlProgram, assetID *bc.AssetID) *account.UTXO {
        utxo := &account.UTXO{}
        utxo.OutputID = bc.Hash{V0: 1}
        utxo.SourceID = bc.Hash{V0: 2}
-       utxo.AssetID = *consensus.BTMAssetID
+       utxo.AssetID = *assetID
        utxo.Amount = 1000000000
        utxo.SourcePos = 0
        utxo.ControlProgram = controlProg.ControlProgram
@@ -210,27 +115,37 @@ func mockUTXO(controlProg *account.CtrlProgram) *account.UTXO {
        return utxo
 }
 
-func mockTxData(utxo *account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
-       txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
-       if err != nil {
-               return nil, nil, err
+func mockTxData(utxos []*account.UTXO, testAccount *account.Account) (*txbuilder.Template, *types.TxData, error) {
+       tplBuilder := txbuilder.NewBuilder(time.Now())
+
+       for _, utxo := range utxos {
+               txInput, sigInst, err := account.UtxoToInputs(testAccount.Signer, utxo)
+               if err != nil {
+                       return nil, nil, err
+               }
+               tplBuilder.AddInput(txInput, sigInst)
+
+               out := &types.TxOutput{}
+               if utxo.AssetID == *consensus.BTMAssetID {
+                       out = types.NewTxOutput(utxo.AssetID, 100, utxo.ControlProgram)
+               } else {
+                       out = types.NewTxOutput(utxo.AssetID, utxo.Amount, utxo.ControlProgram)
+               }
+               tplBuilder.AddOutput(out)
        }
 
-       b := txbuilder.NewBuilder(time.Now())
-       b.AddInput(txInput, sigInst)
-       out := types.NewTxOutput(*consensus.BTMAssetID, 100, utxo.ControlProgram)
-       b.AddOutput(out)
-       return b.Build()
+       return tplBuilder.Build()
 }
 
 func mockWallet(walletDB dbm.DB, account *account.Manager, asset *asset.Registry, chain *protocol.Chain) *Wallet {
-       return &Wallet{
-               DB:             walletDB,
-               AccountMgr:     account,
-               AssetReg:       asset,
-               chain:          chain,
-               rescanProgress: make(chan struct{}, 1),
-       }
+       wallet := &Wallet{
+               DB:          walletDB,
+               AccountMgr:  account,
+               AssetReg:    asset,
+               chain:       chain,
+               RecoveryMgr: newRecoveryManager(walletDB, account),
+       }
+       return wallet
 }
 
 func mockSingleBlock(tx *types.Tx) *types.Block {