message: Create a separate class for the debug mode unit tests to set up and clean up after the tests are run.
diff: @@ -588,9 +588,19 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  "DatasetSpec(Foo(a=TensorSpec(shape=(), dtype=tf.int32, name=None), "
  "b=TensorSpec(shape=(), dtype=tf.string, name=None)), TensorShape([]))")
  +
  +class DebugDatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  +
  + def setUp(self):
  + super(DebugDatasetTest, self).setUp()
  + dataset_ops.toggle_debug_mode(True)
  +
  + def tearDown(self):
  + dataset_ops.toggle_debug_mode(False)
  + super(DebugDatasetTest, self).tearDown()
  +
  @combinations.generate(test_base.eager_only_combinations())
  def testDebugModeEagerExecution(self):
  - dataset_ops.toggle_debug_mode(True)
  counter = []
  ds = dataset_ops.Dataset.range(10)
  @@ -605,11 +615,9 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  # traces the function to figure out what the types and shapes of its
  # outputs are.
  self.assertLen(counter, 11)
  - dataset_ops.toggle_debug_mode(False)
  @combinations.generate(test_base.eager_only_combinations())
  def testDebugModeSequentialExecution(self):
  - dataset_ops.toggle_debug_mode(True)
  ds = dataset_ops.Dataset.range(10)
  ds = ds.apply(
  testing.assert_next(["Interleave", "Map", "Batch", "FiniteTake"]))
  @@ -622,7 +630,6 @@ class DatasetTest(test_base.DatasetTestBase, parameterized.TestCase):
  ds = ds.prefetch(buffer_size=2)
  ds = ds.take(2)
  self.assertDatasetProduces(ds, [[0, 1, 4, 9, 16], [25, 36, 49, 64, 81]])
  - dataset_ops.toggle_debug_mode(False)
  if __name__ == "__main__":
  