def test_runs(self):
    """Test that data mocking utility runs for cases expected to succeed."""

    with TemporaryDirectory() as tmp_dir:

      mock_raw_data(tmp_dir, raw_dim=256, num_channels=3, num_images=40)
Test that data mocking utility runs for cases expected to succeed.
tensor2tensor/data_generators/allen_brain_test.py