test / tests /test_iterator_pipe.py
iblfe's picture
Upload folder using huggingface_hub
b585c7f verified
import unittest
from iterators import IteratorPipe
class TestQueueToIterator(unittest.TestCase):
def test_normal_iteration(self):
it = IteratorPipe()
it.put(1)
it.put(2)
it.put(3)
it.close() # stop iteration
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_normal_custom_sentinel(self):
sentinel = object()
it = IteratorPipe(sentinel=sentinel)
it.put(1)
it.put(2)
it.put(3)
it.put(sentinel) # stop iteration
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_multiple_close(self):
sentinel = object()
it = IteratorPipe(sentinel=sentinel)
it.put(1)
it.put(2)
it.put(3)
it.close() # stop iteration
it.close() # stop iteration
it.close() # stop iteration
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_multiple_next_after_close(self):
sentinel = object()
it = IteratorPipe(sentinel=sentinel)
it.put(1)
it.put(2)
it.put(3)
it.close() # stop iteration
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
self.assertRaises(StopIteration, next, it)
self.assertRaises(StopIteration, next, it)
def test_put_after_close(self):
sentinel = object()
it = IteratorPipe(sentinel=sentinel)
self.assertTrue(it.put(1))
it.close() # stop iteration
self.assertFalse(it.put(2))
it.close() # stop iteration
self.assertFalse(it.put(3))
it.close() # stop iteration
self.assertEqual(next(it), 1)
self.assertRaises(StopIteration, next, it)
self.assertRaises(StopIteration, next, it)