github.com/cilium/statedb@v0.3.2/iterator_test.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package statedb
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/cilium/statedb/index"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestCollectFilterMapToSeq(t *testing.T) {
    15  	type testObject struct {
    16  		ID int
    17  	}
    18  
    19  	db := New()
    20  	idIndex := Index[*testObject, int]{
    21  		Name: "id",
    22  		FromObject: func(t *testObject) index.KeySet {
    23  			return index.NewKeySet(index.Int(t.ID))
    24  		},
    25  		FromKey: index.Int,
    26  		Unique:  true,
    27  	}
    28  	table, _ := NewTable("test", idIndex)
    29  	require.NoError(t, db.RegisterTable(table))
    30  	db.Start()
    31  	defer db.Stop()
    32  
    33  	txn := db.WriteTxn(table)
    34  	table.Insert(txn, &testObject{ID: 1})
    35  	table.Insert(txn, &testObject{ID: 2})
    36  	table.Insert(txn, &testObject{ID: 3})
    37  	table.Insert(txn, &testObject{ID: 4})
    38  	table.Insert(txn, &testObject{ID: 5})
    39  	txn.Commit()
    40  
    41  	iter := table.All(db.ReadTxn())
    42  	filtered := Collect(
    43  		Map(
    44  			Filter(
    45  				iter,
    46  				func(obj *testObject) bool {
    47  					return obj.ID%2 == 0
    48  				},
    49  			),
    50  			func(obj *testObject) int {
    51  				return obj.ID
    52  			},
    53  		),
    54  	)
    55  	assert.Len(t, filtered, 2)
    56  	assert.Equal(t, []int{2, 4}, filtered)
    57  
    58  	count := 0
    59  	for obj := range ToSeq(iter) {
    60  		assert.Greater(t, obj.ID, 0)
    61  		count++
    62  	}
    63  	assert.Equal(t, 5, count)
    64  
    65  }