github.com/ava-labs/avalanchego@v1.11.11/vms/platformvm/state/staker_diff_iterator_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package state
     5  
     6  import (
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/ava-labs/avalanchego/ids"
    13  	"github.com/ava-labs/avalanchego/utils/iterator"
    14  	"github.com/ava-labs/avalanchego/vms/platformvm/txs"
    15  )
    16  
    17  func TestStakerDiffIterator(t *testing.T) {
    18  	require := require.New(t)
    19  	currentStakers := []*Staker{
    20  		{
    21  			TxID:     ids.GenerateTestID(),
    22  			EndTime:  time.Unix(10, 0),
    23  			NextTime: time.Unix(10, 0),
    24  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
    25  		},
    26  	}
    27  	pendingStakers := []*Staker{
    28  		{
    29  			TxID:      ids.GenerateTestID(),
    30  			StartTime: time.Unix(0, 0),
    31  			EndTime:   time.Unix(5, 0),
    32  			NextTime:  time.Unix(0, 0),
    33  			Priority:  txs.PrimaryNetworkDelegatorApricotPendingPriority,
    34  		},
    35  		{
    36  			TxID:      ids.GenerateTestID(),
    37  			StartTime: time.Unix(5, 0),
    38  			EndTime:   time.Unix(10, 0),
    39  			NextTime:  time.Unix(5, 0),
    40  			Priority:  txs.PrimaryNetworkDelegatorApricotPendingPriority,
    41  		},
    42  		{
    43  			TxID:      ids.GenerateTestID(),
    44  			StartTime: time.Unix(11, 0),
    45  			EndTime:   time.Unix(20, 0),
    46  			NextTime:  time.Unix(11, 0),
    47  			Priority:  txs.PrimaryNetworkValidatorPendingPriority,
    48  		},
    49  		{
    50  			TxID:      ids.GenerateTestID(),
    51  			StartTime: time.Unix(11, 0),
    52  			EndTime:   time.Unix(20, 0),
    53  			NextTime:  time.Unix(11, 0),
    54  			Priority:  txs.PrimaryNetworkDelegatorApricotPendingPriority,
    55  		},
    56  	}
    57  
    58  	stakerDiffs := []struct {
    59  		txID    ids.ID
    60  		isAdded bool
    61  	}{
    62  		{
    63  			txID:    pendingStakers[0].TxID,
    64  			isAdded: true,
    65  		},
    66  		{
    67  			txID:    pendingStakers[1].TxID,
    68  			isAdded: true,
    69  		},
    70  		{
    71  			txID:    pendingStakers[0].TxID,
    72  			isAdded: false,
    73  		},
    74  		{
    75  			txID:    pendingStakers[1].TxID,
    76  			isAdded: false,
    77  		},
    78  		{
    79  			txID:    currentStakers[0].TxID,
    80  			isAdded: false,
    81  		},
    82  		{
    83  			txID:    pendingStakers[2].TxID,
    84  			isAdded: true,
    85  		},
    86  		{
    87  			txID:    pendingStakers[3].TxID,
    88  			isAdded: true,
    89  		},
    90  		{
    91  			txID:    pendingStakers[3].TxID,
    92  			isAdded: false,
    93  		},
    94  		{
    95  			txID:    pendingStakers[2].TxID,
    96  			isAdded: false,
    97  		},
    98  	}
    99  
   100  	it := NewStakerDiffIterator(
   101  		iterator.FromSlice(currentStakers...),
   102  		iterator.FromSlice(pendingStakers...),
   103  	)
   104  	for _, expectedStaker := range stakerDiffs {
   105  		require.True(it.Next())
   106  		staker, isAdded := it.Value()
   107  		require.Equal(expectedStaker.txID, staker.TxID)
   108  		require.Equal(expectedStaker.isAdded, isAdded)
   109  	}
   110  	require.False(it.Next())
   111  	it.Release()
   112  	require.False(it.Next())
   113  }
   114  
   115  func TestMutableStakerIterator(t *testing.T) {
   116  	require := require.New(t)
   117  	initialStakers := []*Staker{
   118  		{
   119  			TxID:     ids.GenerateTestID(),
   120  			EndTime:  time.Unix(10, 0),
   121  			NextTime: time.Unix(10, 0),
   122  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   123  		},
   124  		{
   125  			TxID:     ids.GenerateTestID(),
   126  			EndTime:  time.Unix(20, 0),
   127  			NextTime: time.Unix(20, 0),
   128  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   129  		},
   130  		{
   131  			TxID:     ids.GenerateTestID(),
   132  			EndTime:  time.Unix(30, 0),
   133  			NextTime: time.Unix(30, 0),
   134  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   135  		},
   136  	}
   137  
   138  	it := newMutableStakerIterator(iterator.FromSlice(initialStakers...))
   139  
   140  	addedStakers := []*Staker{
   141  		{
   142  			TxID:     ids.GenerateTestID(),
   143  			EndTime:  time.Unix(5, 0),
   144  			NextTime: time.Unix(5, 0),
   145  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   146  		},
   147  		{
   148  			TxID:     ids.GenerateTestID(),
   149  			EndTime:  time.Unix(15, 0),
   150  			NextTime: time.Unix(15, 0),
   151  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   152  		},
   153  		{
   154  			TxID:     ids.GenerateTestID(),
   155  			EndTime:  time.Unix(25, 0),
   156  			NextTime: time.Unix(25, 0),
   157  			Priority: txs.PrimaryNetworkValidatorCurrentPriority,
   158  		},
   159  	}
   160  
   161  	// Next must be called before Add.
   162  	hasNext := it.Next()
   163  	require.True(hasNext)
   164  
   165  	for _, staker := range addedStakers {
   166  		it.Add(staker)
   167  	}
   168  
   169  	stakerDiffs := []ids.ID{
   170  		addedStakers[0].TxID,
   171  		initialStakers[0].TxID,
   172  		addedStakers[1].TxID,
   173  		initialStakers[1].TxID,
   174  		addedStakers[2].TxID,
   175  		initialStakers[2].TxID,
   176  	}
   177  
   178  	for _, expectedStakerTxID := range stakerDiffs {
   179  		require.True(hasNext)
   180  		staker := it.Value()
   181  		require.Equal(expectedStakerTxID, staker.TxID)
   182  		hasNext = it.Next()
   183  	}
   184  	require.False(hasNext)
   185  	it.Release()
   186  	require.False(it.Next())
   187  }