scheduler: Adding SetLimit to LimitIterator

This commit is contained in:
Armon Dadgar 2015-09-11 12:01:22 -07:00
parent 5afa7e4681
commit e804567324
2 changed files with 6 additions and 1 deletions

View File

@ -19,6 +19,10 @@ func NewLimitIterator(ctx Context, source RankIterator, limit int) *LimitIterato
return iter
}
func (iter *LimitIterator) SetLimit(limit int) {
iter.limit = limit
}
func (iter *LimitIterator) Next() *RankedNode {
if iter.seen == iter.limit {
return nil

View File

@ -24,7 +24,8 @@ func TestLimitIterator(t *testing.T) {
}
static := NewStaticRankIterator(ctx, nodes)
limit := NewLimitIterator(ctx, static, 2)
limit := NewLimitIterator(ctx, static, 1)
limit.SetLimit(2)
out := collectRanked(limit)
if len(out) != 2 {