def test_take_peeked(self) -> None:
it = PeekingIterator(range(2))
self.assertEqual(0, it.peek())
it.take_peeked(0)
self.assertEqual(1, next(it))
with self.assertRaises(StopIteration):
next(it)
def test_peek_after_exhaustion(self) -> None:
it = PeekingIterator(range(2))
self.assertEqual(0, next(it))
self.assertEqual(1, next(it))
with self.assertRaises(StopIteration):
next(it)
with self.assertRaises(StopIteration):
it.peek()
self.assertEqual(999, it.peek_default(999))
def test_no_peek(self) -> None:
it = PeekingIterator(range(3))
self.assertEqual(0, next(it))
self.assertEqual(1, next(it))
self.assertEqual(2, next(it))
with self.assertRaises(StopIteration):
next(it)
def test_take_peeked_wrong_value(self) -> None:
it = PeekingIterator(range(2))
self.assertEqual(0, it.peek())
with self.assertRaisesRegex(
AssertionError, 'expected the peaked value to be the same'):
it.take_peeked(1)
it.take_peeked(0)
self.assertEqual(1, next(it))
def test_one_chunk_logging(self) -> None:
it = PeekingIterator(range(1, 4))
actual, has_more = one_chunk(it=it, n=2, metric=lambda x: x)
self.assertSequenceEqual([1], tuple(actual))
self.assertTrue(has_more)
actual, has_more = one_chunk(it=it, n=2, metric=lambda x: x)
self.assertSequenceEqual([2], tuple(actual))
self.assertTrue(has_more)
with self.assertLogs(logger='amundsen_gremlin.utils.streams',
level=logging.ERROR):
actual, has_more = one_chunk(it=it, n=2, metric=lambda x: x)
self.assertSequenceEqual([3], tuple(actual))
self.assertFalse(has_more)
def test_peek_is_next(self) -> None:
it = PeekingIterator(range(2))
self.assertEqual(0, it.peek())
self.assertTrue(it.has_more())
self.assertEqual(0, next(it))
self.assertTrue(it.has_more())
self.assertEqual(1, next(it))
self.assertFalse(it.has_more())
with self.assertRaises(StopIteration):
next(it)
def test_peek_repeats(self) -> None:
it = PeekingIterator(range(2))
for _ in range(100):
self.assertEqual(0, it.peek())
self.assertEqual(0, next(it))
self.assertEqual(1, next(it))