def test_with_dataset_repeat_after_cache(synthetic_dataset, reader_factory):
    """ Check if ``tf.data.Dataset``'s ``repeat`` works after ``tf.data.Dataset``'s ``cache``."""
    epochs = 3
    with reader_factory(synthetic_dataset.url, schema_fields=[TestSchema.id]) as reader:
        dataset = make_petastorm_dataset(reader)
        dataset = dataset.cache()
        dataset = dataset.repeat(epochs)
        iterator = dataset.make_one_shot_iterator()
        it_op = iterator.get_next()
        # Check if dataset generates same result in every epoch.
        with tf.Session() as sess:
            with pytest.warns(None):
                # Expect no warnings since cache() is called before repeat()
                for _ in range(epochs):
                    actual_res = []
                    for _, _ in enumerate(synthetic_dataset.data):
                        actual = sess.run(it_op)._asdict()
                        actual_res.append(actual["id"])
                    expected_res = list(range(len(synthetic_dataset.data)))
                    # sort dataset output since row_groups are shuffled from reader.
                    np.testing.assert_equal(sorted(actual_res), expected_res)

            # Exhausted all epochs. Fetching next value should trigger OutOfRangeError
            with pytest.raises(tf.errors.OutOfRangeError):
                sess.run(it_op)
 Check if ``tf.data.Dataset``'s ``repeat`` works after ``tf.data.Dataset``'s ``cache``.
petastorm/tests/test_tf_dataset.py