tests/test_iterator_pipe.py

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)