diff --git a/core/fx/fn.go b/core/fx/fn.go index 53be667d..7c802048 100644 --- a/core/fx/fn.go +++ b/core/fx/fn.go @@ -84,6 +84,14 @@ func (p Stream) Buffer(n int) Stream { return Range(source) } +// Count counts the number of elements in the result. +func (p Stream) Count() (count int) { + for range p.source { + count++ + } + return +} + // Distinct removes the duplicated items base on the given KeyFunc. func (p Stream) Distinct(fn KeyFunc) Stream { source := make(chan interface{}) diff --git a/core/fx/fn_test.go b/core/fx/fn_test.go index 8ff38ec8..a1a2ac6c 100644 --- a/core/fx/fn_test.go +++ b/core/fx/fn_test.go @@ -49,6 +49,36 @@ func TestBufferNegative(t *testing.T) { assert.Equal(t, 10, result) } +func TestCount(t *testing.T) { + tests := []struct { + name string + elements []interface{} + }{ + { + name: "no elements with nil", + }, + { + name: "no elements", + elements: []interface{}{}, + }, + { + name: "1 element", + elements: []interface{}{1}, + }, + { + name: "multiple elements", + elements: []interface{}{1, 2, 3}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + val := Just(test.elements...).Count() + assert.Equal(t, len(test.elements), val) + }) + } +} + func TestDone(t *testing.T) { var count int32 Just(1, 2, 3).Walk(func(item interface{}, pipe chan<- interface{}) {