Parameterize the tests in function_test.py
# See the License for the specific language governing permissions and
# limitations under the License.
+from absl.testing import parameterized
import tensorflow as tf
from tensorflow_federated.python.tensorflow_libs import function
-class FunctionTest(tf.test.TestCase):
+class FunctionTest(tf.test.TestCase, parameterized.TestCase):
- def test_is_tf_function(self):
- self.assertTrue(function.is_tf_function(tf.function(lambda x: None)))
- concrete_fn = tf.function(
- lambda x: None, input_signature=(tf.TensorSpec(None, tf.int32),))
- self.assertTrue(function.is_tf_function(concrete_fn))
- self.assertFalse(function.is_tf_function(lambda x: None))
- self.assertFalse(function.is_tf_function(None))
+ @parameterized.named_parameters(
+ ('polymorphic_function', tf.function(lambda x: None)),
+ ('concrete_function',
+ tf.function(
+ lambda x: None, input_signature=(tf.TensorSpec(None, tf.int32),))),
+ )
+ def test_is_tf_function_true(self, fn):
+ self.assertTrue(function.is_tf_function(fn))
+
+ @parameterized.named_parameters(
+ ('lambda', lambda x: None),
+ ('none', None),
+ )
+ def test_is_tf_function_false(self, fn):
+ self.assertFalse(function.is_tf_function(fn))
if __name__ == '__main__':
