diff options
Diffstat (limited to 'tensorflow/python/platform')
-rw-r--r-- | tensorflow/python/platform/googletest.py | 2 | ||||
-rw-r--r-- | tensorflow/python/platform/test.py | 2 |
2 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/python/platform/googletest.py b/tensorflow/python/platform/googletest.py index 96219faab7..8141cf92c5 100644 --- a/tensorflow/python/platform/googletest.py +++ b/tensorflow/python/platform/googletest.py @@ -36,6 +36,7 @@ from tensorflow.python.platform import benchmark from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect +from tensorflow.python.util.tf_export import tf_export Benchmark = benchmark.TensorFlowBenchmark # pylint: disable=invalid-name @@ -138,6 +139,7 @@ def StatefulSessionAvailable(): return False +@tf_export('test.StubOutForTesting') class StubOutForTesting(object): """Support class for stubbing methods out for unit testing. diff --git a/tensorflow/python/platform/test.py b/tensorflow/python/platform/test.py index 9b7655722a..1660791feb 100644 --- a/tensorflow/python/platform/test.py +++ b/tensorflow/python/platform/test.py @@ -62,6 +62,8 @@ if sys.version_info.major == 2: else: from unittest import mock # pylint: disable=g-import-not-at-top +tf_export('test.mock')(mock) + # Import Benchmark class Benchmark = _googletest.Benchmark # pylint: disable=invalid-name |