def test_simple_ngram_read_tensorflow(synthetic_dataset):
    """Read a single ngram. Make sure all shapes are set and the data read matches reference data"""
    fields = {
        0: [TestSchema.id],
        1: [TestSchema.id],
        2: [TestSchema.id]
    }

    # Expecting delta between ids to be 1. Setting 1.5 as upper bound
    ngram = NGram(fields=fields, delta_threshold=1.5, timestamp_field=TestSchema.id)

    ngrams, row_tensors_seq = \
        _read_from_tf_tensors(synthetic_dataset, 30, shuffling_queue_capacity=0, min_after_dequeue=0, ngram=ngram)

    for row_tensors in row_tensors_seq.values():
        _assert_all_tensors_have_shape(row_tensors)

    for one_ngram_dict in ngrams:
        _assert_expected_rows_data(synthetic_dataset.data, one_ngram_dict.values())
Read a single ngram. Make sure all shapes are set and the data read matches reference data
petastorm/tests/test_tf_utils.py