async def test_tee_list_six(self): itrs = ait.tee(slist * 2, n=6) for k in slist * 2: values = await asyncio.gather(*[ait.next(it) for it in itrs]) for value in values: self.assertEqual(value, k) for it in itrs: with self.assertRaises(StopAsyncIteration): await ait.next(it)
async def test_tee_list_two(self): it1, it2 = ait.tee(slist * 2) for k in slist * 2: a, b = await asyncio.gather(ait.next(it1), ait.next(it2)) self.assertEqual(a, b) self.assertEqual(a, k) self.assertEqual(b, k) for it in [it1, it2]: with self.assertRaises(StopAsyncIteration): await ait.next(it)
async def test_tee_gen_six(self): async def gen(): yield 1 yield 4 yield 9 yield 16 itrs = ait.tee(gen(), n=6) for k in [1, 4, 9, 16]: values = await asyncio.gather(*[ait.next(it) for it in itrs]) for value in values: self.assertEqual(value, k) for it in itrs: with self.assertRaises(StopAsyncIteration): await ait.next(it)
async def test_tee_gen_two(self): async def gen(): yield 1 yield 4 yield 9 yield 16 it1, it2 = ait.tee(gen()) for k in [1, 4, 9, 16]: a, b = await asyncio.gather(ait.next(it1), ait.next(it2)) self.assertEqual(a, b) self.assertEqual(a, k) self.assertEqual(b, k) for it in [it1, it2]: with self.assertRaises(StopAsyncIteration): await ait.next(it)
def result_key_iters(self): teed_results = aioitertools.tee(self, len(self.result_keys)) return [ ResultKeyIterator(i, result_key) for i, result_key in zip(teed_results, self.result_keys) ]