test / tests /test_timeout_iterator.py
iblfe's picture
Upload folder using huggingface_hub
b585c7f verified
import unittest
import time
from iterators import TimeoutIterator
def iter_simple():
yield 1
yield 2
def iter_with_sleep():
yield 1
time.sleep(0.6)
yield 2
time.sleep(0.4)
yield 3
def iter_with_exception():
yield 1
yield 2
raise Exception
yield 3
class TestTimeoutIterator(unittest.TestCase):
def test_normal_iteration(self):
i = iter_simple()
it = TimeoutIterator(i)
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertRaises(StopIteration, next, it)
self.assertRaises(StopIteration, next, it)
def test_normal_iteration_for_loop(self):
i = iter_simple()
it = TimeoutIterator(i)
iterResults = []
for x in it:
iterResults.append(x)
self.assertEqual(iterResults, [1, 2])
def test_timeout_block(self):
i = iter_with_sleep()
it = TimeoutIterator(i)
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)
def test_timeout_block_for_loop(self):
i = iter_with_sleep()
it = TimeoutIterator(i)
iterResults = []
for x in it:
iterResults.append(x)
self.assertEqual(iterResults, [1, 2, 3])
def test_fixed_timeout(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5)
self.assertEqual(next(it), 1)
self.assertEqual(next(it), it.get_sentinel())
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_fixed_timeout_for_loop(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5)
iterResults = []
for x in it:
iterResults.append(x)
self.assertEqual(iterResults, [1, it.get_sentinel(), 2, 3])
def test_timeout_update(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5)
self.assertEqual(next(it), 1)
self.assertEqual(next(it), it.get_sentinel())
it.set_timeout(0.3)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), it.get_sentinel())
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_custom_sentinel(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5, sentinel="END")
self.assertEqual(next(it), 1)
self.assertEqual(next(it), "END")
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_feature_timeout_reset(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5, reset_on_next=True)
self.assertEqual(next(it), 1) # timeout gets reset after first iteration
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_function_set_reset_on_next(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.35, reset_on_next=False)
self.assertEqual(next(it), 1)
self.assertEqual(next(it), it.get_sentinel())
it.set_reset_on_next(True)
self.assertEqual(next(it), 2)
self.assertEqual(next(it), 3)
self.assertRaises(StopIteration, next, it)
def test_iterator_raises_exception(self):
i = iter_with_exception()
it = TimeoutIterator(i, timeout=0.5, sentinel="END")
self.assertEqual(next(it), 1)
self.assertEqual(next(it), 2)
self.assertRaises(Exception, next, it)
self.assertRaises(StopIteration, next, it)
def test_interrupt_thread(self):
i = iter_with_sleep()
it = TimeoutIterator(i, timeout=0.5, sentinel="END")
self.assertEqual(next(it), 1)
self.assertEqual(next(it), it.get_sentinel())
it.interrupt()
self.assertEqual(next(it), 2)
self.assertRaises(StopIteration, next, it)