1 files changed, 103 insertions, 0 deletions
diff --git a/tensorflow/python/autograph/utils/py_func_test.py b/tensorflow/python/autograph/utils/py_func_test.py
new file mode 100644
index 0000000000..1c220d9492
--- /dev/null
+++ b/tensorflow/python/autograph/utils/py_func_test.py
@@ -0,0 +1,103 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for wrap_py_func module."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.autograph.utils import py_func
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+
+
+class PyFuncTest(test.TestCase):
+
+  def test_wrap_py_func_simple(self):
+
+    def test_fn(a, b, c):
+      return a + b + c
+
+    with self.cached_session() as sess:
+      result = py_func.wrap_py_func(test_fn, dtypes.int64,
+                                    (1, constant_op.constant(1), 1))
+      self.assertEqual(3, sess.run(result))
+      result = py_func.wrap_py_func(test_fn, dtypes.int64, (1, 1, 1))
+      self.assertEqual(3, sess.run(result))
+      result = py_func.wrap_py_func(
+          test_fn, dtypes.int64,
+          (constant_op.constant(1), 1, constant_op.constant(1)))
+      self.assertEqual(3, sess.run(result))
+
+  def test_wrap_py_func_complex_args(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.foo = 5
+
+    def test_fn(a, b):
+      return a * b.foo
+
+    with self.cached_session() as sess:
+      result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass()))
+      self.assertEqual(35, sess.run(result))
+      result = py_func.wrap_py_func(test_fn, dtypes.int64,
+                                    (constant_op.constant(7), TestClass()))
+      self.assertEqual(35, sess.run(result))
+
+  def test_wrap_py_func_kwargs(self):
+
+    class TestClass(object):
+
+      def __init__(self, foo):
+        self.foo = foo
+
+    def test_fn(a, b, c, d):
+      return a * b.foo + c * d.foo
+
+    with self.cached_session() as sess:
+      result = py_func.wrap_py_func(test_fn, dtypes.int64, (7, TestClass(5)), {
+          'c': 11,
+          'd': TestClass(13)
+      })
+      self.assertEqual(178, sess.run(result))
+      result = py_func.wrap_py_func(test_fn, dtypes.int64,
+                                    (constant_op.constant(7), TestClass(5)), {
+                                        'c': constant_op.constant(11),
+                                        'd': TestClass(13)
+                                    })
+      self.assertEqual(178, sess.run(result))
+
+  def test_wrap_py_func_dummy_return(self):
+
+    side_counter = [0]
+
+    def test_fn(_):
+      side_counter[0] += 1
+
+    with self.cached_session() as sess:
+      result = py_func.wrap_py_func(test_fn, None, (5,), use_dummy_return=True)
+      self.assertEqual(1, sess.run(result))
+      self.assertEqual([1], side_counter)
+      result = py_func.wrap_py_func(
+          test_fn, None, (constant_op.constant(5),), use_dummy_return=True)
+      self.assertEqual(1, sess.run(result))
+      self.assertEqual([2], side_counter)
+
+
+if __name__ == '__main__':
+  test.main()