File size: 1,370 Bytes
fc3399e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
from bytelatent.data.iterators.dev_iterators import BltTestIterator
from bytelatent.data.iterators.limit_iterator import LimitIterator


def test_limit_iterator():
    total = 10
    limit = 5
    base_iterator = BltTestIterator(total=total)
    limit_iterator = LimitIterator(base_iterator, limit=limit)
    iterator = limit_iterator.create_iter()
    n = 0
    for example in iterator:
        assert example.sample_id == f"test_{n}"
        n += 1
    assert n == limit

    limit = 10
    base_iterator = BltTestIterator(total=total)
    limit_iterator = LimitIterator(base_iterator, limit=limit)
    iterator = limit_iterator.create_iter()
    n = 0
    for example in iterator:
        assert example.sample_id == f"test_{n}"
        n += 1
    assert n == limit == total

    limit = 20
    base_iterator = BltTestIterator(total=total)
    limit_iterator = LimitIterator(base_iterator, limit=limit)
    iterator = limit_iterator.create_iter()
    n = 0
    for example in iterator:
        assert example.sample_id == f"test_{n}"
        n += 1
    assert n == total

    limit = -1
    base_iterator = BltTestIterator(total=total)
    limit_iterator = LimitIterator(base_iterator, limit=limit)
    iterator = limit_iterator.create_iter()
    n = 0
    for example in iterator:
        assert example.sample_id == f"test_{n}"
        n += 1
    assert n == total