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

     1  package statesync
     2  
     3  import (
     4  	"encoding/hex"
     5  	"testing"
     6  
     7  	"github.com/nspcc-dev/neo-go/internal/random"
     8  	"github.com/nspcc-dev/neo-go/pkg/util"
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestPool_AddRemoveUpdate(t *testing.T) {
    13  	mp := NewPool()
    14  
    15  	i1 := []byte{1, 2, 3}
    16  	i1h := util.Uint256{1, 2, 3}
    17  	i2 := []byte{2, 3, 4}
    18  	i2h := util.Uint256{2, 3, 4}
    19  	i3 := []byte{4, 5, 6}
    20  	i3h := util.Uint256{3, 4, 5}
    21  	i4 := []byte{3, 4, 5} // has the same hash as i3
    22  	i5 := []byte{6, 7, 8} // has the same hash as i3
    23  	mapAll := map[util.Uint256][][]byte{i1h: {i1}, i2h: {i2}, i3h: {i4, i3}}
    24  
    25  	// No items
    26  	_, ok := mp.TryGet(i1h)
    27  	require.False(t, ok)
    28  	require.False(t, mp.ContainsKey(i1h))
    29  	require.Equal(t, 0, mp.Count())
    30  	require.Equal(t, map[util.Uint256][][]byte{}, mp.GetAll())
    31  
    32  	// Add i1, i2, check OK
    33  	mp.Add(i1h, i1)
    34  	mp.Add(i2h, i2)
    35  	itm, ok := mp.TryGet(i1h)
    36  	require.True(t, ok)
    37  	require.Equal(t, [][]byte{i1}, itm)
    38  	require.True(t, mp.ContainsKey(i1h))
    39  	require.True(t, mp.ContainsKey(i2h))
    40  	require.Equal(t, map[util.Uint256][][]byte{i1h: {i1}, i2h: {i2}}, mp.GetAll())
    41  	require.Equal(t, 2, mp.Count())
    42  
    43  	// Remove i1 and unexisting item
    44  	mp.Remove(i3h)
    45  	mp.Remove(i1h)
    46  	require.False(t, mp.ContainsKey(i1h))
    47  	require.True(t, mp.ContainsKey(i2h))
    48  	require.Equal(t, map[util.Uint256][][]byte{i2h: {i2}}, mp.GetAll())
    49  	require.Equal(t, 1, mp.Count())
    50  
    51  	// Update: remove nothing, add all
    52  	mp.Update(nil, mapAll)
    53  	require.Equal(t, mapAll, mp.GetAll())
    54  	require.Equal(t, 3, mp.Count())
    55  	// Update: remove all, add all
    56  	mp.Update(mapAll, mapAll)
    57  	require.Equal(t, mapAll, mp.GetAll()) // deletion first, addition after that
    58  	require.Equal(t, 3, mp.Count())
    59  	// Update: remove all, add nothing
    60  	mp.Update(mapAll, nil)
    61  	require.Equal(t, map[util.Uint256][][]byte{}, mp.GetAll())
    62  	require.Equal(t, 0, mp.Count())
    63  	// Update: remove several, add several
    64  	mp.Update(map[util.Uint256][][]byte{i1h: {i1}, i2h: {i2}}, map[util.Uint256][][]byte{i2h: {i2}, i3h: {i3}})
    65  	require.Equal(t, map[util.Uint256][][]byte{i2h: {i2}, i3h: {i3}}, mp.GetAll())
    66  	require.Equal(t, 2, mp.Count())
    67  
    68  	// Update: remove nothing, add several with same hashes
    69  	mp.Update(nil, map[util.Uint256][][]byte{i3h: {i5, i4}}) // should be sorted by the pool
    70  	require.Equal(t, map[util.Uint256][][]byte{i2h: {i2}, i3h: {i4, i3, i5}}, mp.GetAll())
    71  	require.Equal(t, 2, mp.Count())
    72  	// Update: remove several with same hashes, add nothing
    73  	mp.Update(map[util.Uint256][][]byte{i3h: {i5, i4}}, nil)
    74  	require.Equal(t, map[util.Uint256][][]byte{i2h: {i2}, i3h: {i3}}, mp.GetAll())
    75  	require.Equal(t, 2, mp.Count())
    76  	// Update: remove several with same hashes, add several with same hashes
    77  	mp.Update(map[util.Uint256][][]byte{i3h: {i5, i3}}, map[util.Uint256][][]byte{i3h: {i5, i4}})
    78  	require.Equal(t, map[util.Uint256][][]byte{i2h: {i2}, i3h: {i4, i5}}, mp.GetAll())
    79  	require.Equal(t, 2, mp.Count())
    80  }
    81  
    82  func TestPool_GetBatch(t *testing.T) {
    83  	check := func(t *testing.T, limit int, itemsCount int) {
    84  		mp := NewPool()
    85  		for i := 0; i < itemsCount; i++ {
    86  			mp.Add(random.Uint256(), []byte{0x01})
    87  		}
    88  		batch := mp.GetBatch(limit)
    89  		if limit < itemsCount {
    90  			require.Equal(t, limit, len(batch))
    91  		} else {
    92  			require.Equal(t, itemsCount, len(batch))
    93  		}
    94  	}
    95  
    96  	t.Run("limit less than items count", func(t *testing.T) {
    97  		check(t, 5, 6)
    98  	})
    99  	t.Run("limit more than items count", func(t *testing.T) {
   100  		check(t, 6, 5)
   101  	})
   102  	t.Run("items count limit", func(t *testing.T) {
   103  		check(t, 5, 5)
   104  	})
   105  }
   106  
   107  func TestPool_UpdateUsingSliceFromPool(t *testing.T) {
   108  	mp := NewPool()
   109  	p1, _ := hex.DecodeString("0f0a0f0f0f0f0f0f0104020b02080c0a06050e070b050404060206060d07080602030b04040b050e040406030f0708060c05")
   110  	p2, _ := hex.DecodeString("0f0a0f0f0f0f0f0f01040a0b000f04000b03090b02090b0e040f0d0b060d070e0b0b090b0906080602060c0d0f0e0d04070e")
   111  	p3, _ := hex.DecodeString("0f0a0f0f0f0f0f0f01040b010d01080f050f000a0d0e08060c040b050800050904060f050807080a080c07040d0107080007")
   112  	h, _ := util.Uint256DecodeStringBE("57e197679ef031bf2f0b466b20afe3f67ac04dcff80a1dc4d12dd98dd21a2511")
   113  	mp.Add(h, p1)
   114  	mp.Add(h, p2)
   115  	mp.Add(h, p3)
   116  
   117  	toBeRemoved, ok := mp.TryGet(h)
   118  	require.True(t, ok)
   119  
   120  	mp.Update(map[util.Uint256][][]byte{h: toBeRemoved}, nil)
   121  	// test that all items were successfully removed.
   122  	require.Equal(t, 0, len(mp.GetAll()))
   123  }