github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/mempool/subscriptions_test.go (about)

     1  package mempool
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/nspcc-dev/neo-go/pkg/core/mempoolevent"
     8  	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
     9  	"github.com/nspcc-dev/neo-go/pkg/util"
    10  	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestSubscriptions(t *testing.T) {
    15  	t.Run("disabled subscriptions", func(t *testing.T) {
    16  		mp := New(5, 0, false, nil)
    17  		require.Panics(t, func() {
    18  			mp.RunSubscriptions()
    19  		})
    20  		require.Panics(t, func() {
    21  			mp.StopSubscriptions()
    22  		})
    23  	})
    24  
    25  	t.Run("enabled subscriptions", func(t *testing.T) {
    26  		fs := &FeerStub{balance: 100}
    27  		mp := New(2, 0, true, nil)
    28  		mp.RunSubscriptions()
    29  		subChan1 := make(chan mempoolevent.Event, 3)
    30  		subChan2 := make(chan mempoolevent.Event, 3)
    31  		mp.SubscribeForTransactions(subChan1)
    32  		t.Cleanup(mp.StopSubscriptions)
    33  
    34  		txs := make([]*transaction.Transaction, 4)
    35  		for i := range txs {
    36  			txs[i] = transaction.New([]byte{byte(opcode.PUSH1)}, 0)
    37  			txs[i].Nonce = uint32(i)
    38  			txs[i].Signers = []transaction.Signer{{Account: util.Uint160{1, 2, 3}}}
    39  			txs[i].NetworkFee = int64(i)
    40  		}
    41  
    42  		// add tx
    43  		require.NoError(t, mp.Add(txs[0], fs))
    44  		require.Eventually(t, func() bool { return len(subChan1) == 1 }, time.Second, time.Millisecond*100)
    45  		event := <-subChan1
    46  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[0]}, event)
    47  
    48  		// severak subscribers
    49  		mp.SubscribeForTransactions(subChan2)
    50  		require.NoError(t, mp.Add(txs[1], fs))
    51  		require.Eventually(t, func() bool { return len(subChan1) == 1 && len(subChan2) == 1 }, time.Second, time.Millisecond*100)
    52  		event1 := <-subChan1
    53  		event2 := <-subChan2
    54  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[1]}, event1)
    55  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[1]}, event2)
    56  
    57  		// reach capacity
    58  		require.NoError(t, mp.Add(txs[2], &FeerStub{}))
    59  		require.Eventually(t, func() bool { return len(subChan1) == 2 && len(subChan2) == 2 }, time.Second, time.Millisecond*100)
    60  		event1 = <-subChan1
    61  		event2 = <-subChan2
    62  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[0]}, event1)
    63  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[0]}, event2)
    64  		event1 = <-subChan1
    65  		event2 = <-subChan2
    66  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[2]}, event1)
    67  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[2]}, event2)
    68  
    69  		// remove tx
    70  		mp.Remove(txs[1].Hash(), fs)
    71  		require.Eventually(t, func() bool { return len(subChan1) == 1 && len(subChan2) == 1 }, time.Second, time.Millisecond*100)
    72  		event1 = <-subChan1
    73  		event2 = <-subChan2
    74  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[1]}, event1)
    75  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[1]}, event2)
    76  
    77  		// remove stale
    78  		mp.RemoveStale(func(tx *transaction.Transaction) bool {
    79  			return !tx.Hash().Equals(txs[2].Hash())
    80  		}, fs)
    81  		require.Eventually(t, func() bool { return len(subChan1) == 1 && len(subChan2) == 1 }, time.Second, time.Millisecond*100)
    82  		event1 = <-subChan1
    83  		event2 = <-subChan2
    84  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[2]}, event1)
    85  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionRemoved, Tx: txs[2]}, event2)
    86  
    87  		// unsubscribe
    88  		mp.UnsubscribeFromTransactions(subChan1)
    89  		require.NoError(t, mp.Add(txs[3], fs))
    90  		require.Eventually(t, func() bool { return len(subChan2) == 1 }, time.Second, time.Millisecond*100)
    91  		event2 = <-subChan2
    92  		require.Equal(t, 0, len(subChan1))
    93  		require.Equal(t, mempoolevent.Event{Type: mempoolevent.TransactionAdded, Tx: txs[3]}, event2)
    94  	})
    95  }