def test_generator_produces_examples(self):
    """Basic test that the generator produces examples with expected keys."""

    for is_training in [True, False]:
      with TemporaryDirectory() as tmp_dir:
        mock_raw_data(tmp_dir, raw_dim=256, num_images=100)
        for example in allen_brain._generator(tmp_dir, is_training):
          for key in ["image/encoded", "image/format",
                      "image/height", "image/width"]:
            self.assertTrue(key in example.keys())
Basic test that the generator produces examples with expected keys.
tensor2tensor/data_generators/allen_brain_test.py