diff --git a/tensorflow/python/framework/ops_enable_eager_test.py b/tensorflow/python/framework/ops_enable_eager_test.py index 99d06f1c2d4..4da0798dd32 100644 --- a/tensorflow/python/framework/ops_enable_eager_test.py +++ b/tensorflow/python/framework/ops_enable_eager_test.py @@ -23,9 +23,10 @@ from tensorflow.python.framework import ops from tensorflow.python.platform import googletest -class OpsEnableEagerTest(googletest.TestCase): +class OpsEnableAndDisableEagerTest(googletest.TestCase): - def test_enable_eager_execution_multiple_times(self): + def setUp(self): + # test for enable eager test ops.enable_eager_execution() self.assertTrue(context.executing_eagerly()) @@ -33,6 +34,15 @@ class OpsEnableEagerTest(googletest.TestCase): ops.enable_eager_execution() self.assertTrue(context.executing_eagerly()) + def tearDown(self): + # test for disable eager test + ops.disable_eager_execution() + self.assertFalse(context.executing_eagerly()) + + # Calling disable eager execution a second time should not cause an error. + ops.disable_eager_execution() + self.assertFalse(context.executing_eagerly()) + if __name__ == '__main__': googletest.main()