| 664 | # set_event_during_invocation`, because we must consume at least |
| 665 | # one element to start the prefetching. |
| 666 | def do_test_ev(buffer_size): |
| 667 | dataset = dataset_ops.Dataset.range(100).map(_map_fn).prefetch( |
| 668 | buffer_size) |
| 669 | |
| 670 | get_next = self.getNext(dataset) |
| 671 | |
| 672 | event_will_be_set_after_consuming = ( |
| 673 | set_event_during_invocation - buffer_size + 1) |
| 674 | |
| 675 | ev.clear() |
| 676 | for i in range(event_will_be_set_after_consuming): |
| 677 | self.assertFalse(ev.is_set()) |
| 678 | self.assertEqual(i * i, self.evaluate(get_next())) |
| 679 | ev.wait() |
| 680 | for i in range(event_will_be_set_after_consuming, 100): |
| 681 | self.assertEqual(i * i, self.evaluate(get_next())) |
| 682 | with self.assertRaises(errors.OutOfRangeError): |
| 683 | self.evaluate(get_next()) |
| 684 | |
| 685 | for buffer_size in range(1, set_event_during_invocation): |
| 686 | do_test_ev(buffer_size) |