Fix test dependencies and move some py_test to py_strict_test when possible

PiperOrigin-RevId: 370802463

Author: teijeong

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -132,7 +132,7 @@ py_test(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "clustering_centroids_test",
     size = "medium",
     srcs = ["clustering_centroids_test.py"],

tensorflow_model_optimization/python/core/clustering/keras/clustering_centroids_test.py

@@ -16,20 +16,17 @@
 
 from absl.testing import parameterized
 import tensorflow as tf
-import tensorflow.keras.backend as K
 
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
 
-keras = tf.keras
+K = tf.keras.backend
 errors_impl = tf.errors
-layers = keras.layers
-test = tf.test
 
 CentroidInitialization = cluster_config.CentroidInitialization
 
 
-class ClusteringCentroidsTest(test.TestCase, parameterized.TestCase):
+class ClusteringCentroidsTest(tf.test.TestCase, parameterized.TestCase):
   """Unit tests for the clustering_centroids module."""
 
   def setUp(self):
@@ -221,4 +218,4 @@ def testKmeansPlusPlusClusterCentroidsWithSparsityPreservation(
 
 
 if __name__ == "__main__":
-  test.main()
+  tf.test.main()

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -131,7 +131,7 @@ py_strict_library(
     ],
 )
 
-py_test(
+py_strict_test(
     name = "quantize_annotate_test",
     srcs = [
         "quantize_annotate_test.py",

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/prune_preserve/BUILD

@@ -41,6 +41,7 @@ py_test(
     visibility = ["//visibility:private"],
     deps = [
         ":prune_preserve_quantize_registry",
+        # absl/testing:parameterized dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry",
         "//tensorflow_model_optimization/python/core/sparsity/keras:prune_registry",

tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate_test.py

@@ -21,7 +21,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.keras.engine.base_layer import Layer
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_annotate
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_config as quantize_config_mod
 
@@ -54,7 +53,7 @@ def get_config(self):
 
   def testAnnotateLayerCallPassesTraningBoolean(self):
 
-    class MockLayer(Layer):
+    class MockLayer(tf.keras.layers.Layer):
       self.training = None
 
       def call(self, training=None):