import tensorflow as tf


def evaluate(defun=False):
  model = mnist.create_model(data_format())
  dataset = random_dataset()
  if defun:
    model.call = tfe.defun(model.call)
  with tf.device(device()):
    mnist_eager.test(model, dataset)


class MNISTTest(tf.test.TestCase):
  """Run tests for MNIST eager loop."""

  def setUp(self):
    if not keras_utils.is_v2_0():
      tf.compat.v1.enable_v2_behavior()
    super(MNISTTest, self).setUp()

  def test_train(self):
    train(defun=False)

  def test_evaluate(self):
    evaluate(defun=False)

  def test_train_with_defun(self):
    train(defun=True)

  def test_evaluate_with_defun(self):
    evaluate(defun=True)

