def test_image_mock_produces_expected_shape(self):
    """Test that the image mocking utility produces expected shape output."""

    with TemporaryDirectory() as tmp_dir:

      cases = [
          {
              "x_dim": 8,
              "y_dim": 8,
              "num_channels": 3,
              "output_path": "/foo",
              "write_image": True
          }
      ]

      for cid, case in enumerate(cases):
        output_path = os.path.join(tmp_dir, "dummy%s.jpg" % cid)
        img = mock_raw_image(x_dim=case["x_dim"],
                             y_dim=case["y_dim"],
                             num_channels=case["num_channels"],
                             output_path=output_path,
                             write_image=case["write_image"])

        self.assertEqual(img.shape, (case["x_dim"], case["y_dim"],
                                     case["num_channels"]))
        if case["write_image"]:
          self.assertTrue(tf.gfile.Exists(output_path))
Test that the image mocking utility produces expected shape output.
tensor2tensor/data_generators/allen_brain_test.py