github.com/MetalBlockchain/metalgo@v1.11.9/vms/platformvm/state/tree_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/google/btree"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  )
    15  
    16  func TestTreeIterator(t *testing.T) {
    17  	require := require.New(t)
    18  	stakers := []*Staker{
    19  		{
    20  			TxID:     ids.GenerateTestID(),
    21  			NextTime: time.Unix(0, 0),
    22  		},
    23  		{
    24  			TxID:     ids.GenerateTestID(),
    25  			NextTime: time.Unix(1, 0),
    26  		},
    27  		{
    28  			TxID:     ids.GenerateTestID(),
    29  			NextTime: time.Unix(2, 0),
    30  		},
    31  	}
    32  
    33  	tree := btree.NewG(defaultTreeDegree, (*Staker).Less)
    34  	for _, staker := range stakers {
    35  		require.Nil(tree.ReplaceOrInsert(staker))
    36  	}
    37  
    38  	it := NewTreeIterator(tree)
    39  	for _, staker := range stakers {
    40  		require.True(it.Next())
    41  		require.Equal(staker, it.Value())
    42  	}
    43  	require.False(it.Next())
    44  	it.Release()
    45  }
    46  
    47  func TestTreeIteratorNil(t *testing.T) {
    48  	it := NewTreeIterator(nil)
    49  	require.False(t, it.Next())
    50  	it.Release()
    51  }
    52  
    53  func TestTreeIteratorEarlyRelease(t *testing.T) {
    54  	require := require.New(t)
    55  	stakers := []*Staker{
    56  		{
    57  			TxID:     ids.GenerateTestID(),
    58  			NextTime: time.Unix(0, 0),
    59  		},
    60  		{
    61  			TxID:     ids.GenerateTestID(),
    62  			NextTime: time.Unix(1, 0),
    63  		},
    64  		{
    65  			TxID:     ids.GenerateTestID(),
    66  			NextTime: time.Unix(2, 0),
    67  		},
    68  	}
    69  
    70  	tree := btree.NewG(defaultTreeDegree, (*Staker).Less)
    71  	for _, staker := range stakers {
    72  		require.Nil(tree.ReplaceOrInsert(staker))
    73  	}
    74  
    75  	it := NewTreeIterator(tree)
    76  	require.True(it.Next())
    77  	it.Release()
    78  	require.False(it.Next())
    79  }