From 12ec80d2395da75c18db5b8371c68f07889dba11 Mon Sep 17 00:00:00 2001
From: Scott Zhu <scottzhu@google.com>
Date: Tue, 16 Jun 2020 09:18:02 -0700
Subject: [PATCH] Move the module test wrt to Keras private API to keras test.

PiperOrigin-RevId: 316695477
Change-Id: I4201f62e97bb1acb768b980a51d26bf0c76853ab
---
 tensorflow/python/keras/layers/BUILD          | 12 ++++++
 tensorflow/python/keras/layers/layers_test.py | 40 +++++++++++++++++++
 tensorflow/tools/api/tests/module_test.py     | 12 ------
 3 files changed, 52 insertions(+), 12 deletions(-)
 create mode 100644 tensorflow/python/keras/layers/layers_test.py

diff --git a/tensorflow/python/keras/layers/BUILD b/tensorflow/python/keras/layers/BUILD
index 91bace936d3..f67838d591a 100644
--- a/tensorflow/python/keras/layers/BUILD
+++ b/tensorflow/python/keras/layers/BUILD
@@ -886,3 +886,15 @@ tf_py_test(
         "@absl_py//absl/testing:parameterized",
     ],
 )
+
+tf_py_test(
+    name = "layers_test",
+    size = "small",
+    srcs = ["layers_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":layers",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:tf2",
+    ],
+)
diff --git a/tensorflow/python/keras/layers/layers_test.py b/tensorflow/python/keras/layers/layers_test.py
new file mode 100644
index 00000000000..35ba029dd83
--- /dev/null
+++ b/tensorflow/python/keras/layers/layers_test.py
@@ -0,0 +1,40 @@
+# Copyright 2020 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.
+# ==============================================================================
+# pylint: disable=g-classes-have-attributes
+"""Tests for layers.__init__."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python import tf2
+from tensorflow.python.keras import layers
+from tensorflow.python.platform import test
+
+
+class LayersTest(test.TestCase):
+
+  def test_keras_private_symbol(self):
+    normalization_parent = layers.Normalization.__module__.split('.')[-1]
+    if tf2.enabled():
+      self.assertEqual('normalization', normalization_parent)
+      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
+    else:
+      self.assertEqual('normalization_v1', normalization_parent)
+      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
+
+
+if __name__ == '__main__':
+  test.main()
diff --git a/tensorflow/tools/api/tests/module_test.py b/tensorflow/tools/api/tests/module_test.py
index aa8a224d00b..5397278f5f3 100644
--- a/tensorflow/tools/api/tests/module_test.py
+++ b/tensorflow/tools/api/tests/module_test.py
@@ -79,18 +79,6 @@ class ModuleTest(test.TestCase):
       tf.compat.v1.summary.FileWriter
     # pylint: enable=pointless-statement
 
-  def testInternalKerasImport(self):
-    # pylint: disable=g-import-not-at-top
-    from tensorflow.python.keras import layers
-    normalization_parent = layers.Normalization.__module__.split('.')[-1]
-    if tf._major_api_version == 2:
-      self.assertEqual('normalization', normalization_parent)
-      self.assertTrue(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    else:
-      self.assertEqual('normalization_v1', normalization_parent)
-      self.assertFalse(layers.BatchNormalization._USE_V2_BEHAVIOR)
-    # pylint: enable=g-import-not-at-top
-
 
 if __name__ == '__main__':
   test.main()