github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/neatptc/helper_test.go (about)

     1  package neatptc
     2  
     3  import (
     4  	"crypto/ecdsa"
     5  	"crypto/rand"
     6  	"math/big"
     7  	"sort"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/neatlab/neatio/chain/core"
    12  	"github.com/neatlab/neatio/chain/core/rawdb"
    13  	"github.com/neatlab/neatio/chain/core/types"
    14  	"github.com/neatlab/neatio/chain/core/vm"
    15  	"github.com/neatlab/neatio/neatdb"
    16  	"github.com/neatlab/neatio/neatptc/downloader"
    17  	"github.com/neatlab/neatio/network/p2p"
    18  	"github.com/neatlab/neatio/network/p2p/discover"
    19  	"github.com/neatlab/neatio/params"
    20  	"github.com/neatlab/neatio/utilities/common"
    21  	"github.com/neatlab/neatio/utilities/crypto"
    22  	"github.com/neatlab/neatio/utilities/event"
    23  )
    24  
    25  var (
    26  	testBankKey, _ = crypto.HexToECDSA("b71c71a67e1177ad4e901695e1b4b9ee17ae16c6668d313eac2f96dbcda3f291")
    27  	testBank       = crypto.PubkeyToAddress(testBankKey.PublicKey)
    28  )
    29  
    30  func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func(int, *core.BlockGen), newtx chan<- []*types.Transaction) (*ProtocolManager, neatdb.Database, error) {
    31  	var (
    32  		evmux = new(event.TypeMux)
    33  		db    = rawdb.NewMemoryDatabase()
    34  		gspec = &core.Genesis{
    35  			Config: params.TestChainConfig,
    36  			Alloc:  core.GenesisAlloc{testBank: {Balance: big.NewInt(1000000)}},
    37  		}
    38  		genesis       = gspec.MustCommit(db)
    39  		blockchain, _ = core.NewBlockChain(db, nil, gspec.Config, nil, vm.Config{}, nil)
    40  	)
    41  	chain, _ := core.GenerateChain(gspec.Config, genesis, nil, db, blocks, generator)
    42  	if _, err := blockchain.InsertChain(chain); err != nil {
    43  		panic(err)
    44  	}
    45  
    46  	pm, err := NewProtocolManager(gspec.Config, mode, DefaultConfig.NetworkId, evmux, &testTxPool{added: newtx}, engine, blockchain, db, nil)
    47  	if err != nil {
    48  		return nil, nil, err
    49  	}
    50  	pm.Start(1000)
    51  	return pm, db, nil
    52  }
    53  
    54  func newTestProtocolManagerMust(t *testing.T, mode downloader.SyncMode, blocks int, generator func(int, *core.BlockGen), newtx chan<- []*types.Transaction) (*ProtocolManager, neatdb.Database) {
    55  	pm, db, err := newTestProtocolManager(mode, blocks, generator, newtx)
    56  	if err != nil {
    57  		t.Fatalf("Failed to create protocol manager: %v", err)
    58  	}
    59  	return pm, db
    60  }
    61  
    62  type testTxPool struct {
    63  	txFeed event.Feed
    64  	pool   []*types.Transaction
    65  	added  chan<- []*types.Transaction
    66  
    67  	lock sync.RWMutex
    68  }
    69  
    70  func (p *testTxPool) AddRemotes(txs []*types.Transaction) []error {
    71  	p.lock.Lock()
    72  	defer p.lock.Unlock()
    73  
    74  	p.pool = append(p.pool, txs...)
    75  	if p.added != nil {
    76  		p.added <- txs
    77  	}
    78  	return make([]error, len(txs))
    79  }
    80  
    81  func (p *testTxPool) Pending() (map[common.Address]types.Transactions, error) {
    82  	p.lock.RLock()
    83  	defer p.lock.RUnlock()
    84  
    85  	batches := make(map[common.Address]types.Transactions)
    86  	for _, tx := range p.pool {
    87  		from, _ := types.Sender(types.HomesteadSigner{}, tx)
    88  		batches[from] = append(batches[from], tx)
    89  	}
    90  	for _, batch := range batches {
    91  		sort.Sort(types.TxByNonce(batch))
    92  	}
    93  	return batches, nil
    94  }
    95  
    96  func (p *testTxPool) SubscribeTxPreEvent(ch chan<- core.TxPreEvent) event.Subscription {
    97  	return p.txFeed.Subscribe(ch)
    98  }
    99  
   100  func newTestTransaction(from *ecdsa.PrivateKey, nonce uint64, datasize int) *types.Transaction {
   101  	tx := types.NewTransaction(nonce, common.Address{}, big.NewInt(0), 100000, big.NewInt(0), make([]byte, datasize))
   102  	tx, _ = types.SignTx(tx, types.HomesteadSigner{}, from)
   103  	return tx
   104  }
   105  
   106  type testPeer struct {
   107  	net p2p.MsgReadWriter
   108  	app *p2p.MsgPipeRW
   109  	*peer
   110  }
   111  
   112  func newTestPeer(name string, version int, pm *ProtocolManager, shake bool) (*testPeer, <-chan error) {
   113  
   114  	app, net := p2p.MsgPipe()
   115  
   116  	var id discover.NodeID
   117  	rand.Read(id[:])
   118  
   119  	peer := pm.newPeer(version, p2p.NewPeer(id, name, nil), net)
   120  
   121  	errc := make(chan error, 1)
   122  	go func() {
   123  		select {
   124  		case pm.newPeerCh <- peer:
   125  			errc <- pm.handle(peer)
   126  		case <-pm.quitSync:
   127  			errc <- p2p.DiscQuitting
   128  		}
   129  	}()
   130  	tp := &testPeer{app: app, net: net, peer: peer}
   131  
   132  	if shake {
   133  		var (
   134  			genesis = pm.blockchain.Genesis()
   135  			head    = pm.blockchain.CurrentHeader()
   136  			td      = pm.blockchain.GetTd(head.Hash(), head.Number.Uint64())
   137  		)
   138  		tp.handshake(nil, td, head.Hash(), genesis.Hash())
   139  	}
   140  	return tp, errc
   141  }
   142  
   143  func (p *testPeer) handshake(t *testing.T, td *big.Int, head common.Hash, genesis common.Hash) {
   144  	msg := &statusData{
   145  		ProtocolVersion: uint32(p.version),
   146  		NetworkId:       DefaultConfig.NetworkId,
   147  		TD:              td,
   148  		CurrentBlock:    head,
   149  		GenesisBlock:    genesis,
   150  	}
   151  	if err := p2p.ExpectMsg(p.app, StatusMsg, msg); err != nil {
   152  		t.Fatalf("status recv: %v", err)
   153  	}
   154  	if err := p2p.Send(p.app, StatusMsg, msg); err != nil {
   155  		t.Fatalf("status send: %v", err)
   156  	}
   157  }
   158  
   159  func (p *testPeer) close() {
   160  	p.app.Close()
   161  }