def test_multiple_epochs(synthetic_dataset, reader_factory):
    """Tests that multiple epochs works as expected"""
    num_epochs = 5
    with reader_factory(synthetic_dataset.url, num_epochs=num_epochs) as reader:
        # Read all expected entries from the dataset and compare the data to reference
        single_epoch_id_set = [d['id'] for d in synthetic_dataset.data]
        actual_ids_in_all_epochs = _readout_all_ids(reader)
        np.testing.assert_equal(sorted(actual_ids_in_all_epochs), sorted(num_epochs * single_epoch_id_set))

        # Reset reader should reset ventilator. Should produce another `num_epochs` results
        reader.reset()
        actual_ids_in_all_epochs = _readout_all_ids(reader)
        np.testing.assert_equal(sorted(actual_ids_in_all_epochs), sorted(num_epochs * single_epoch_id_set))
Tests that multiple epochs works as expected
petastorm/tests/test_end_to_end.py