Register py_test and tf_py_test macros.

PiperOrigin-RevId: 513914259
GCABC123 · Mar 6, 2023 · 02c6f0d · 02c6f0d
1 parent d7c03de
commit 02c6f0d
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
@@ -2507,6 +2507,11 @@ def py_test(deps = [], data = [], kernels = [], exec_properties = None, **kwargs
         **kwargs
     )
 
+register_extension_info(
+    extension = py_test,
+    label_regex_for_dep = "{extension_name}",
+)
+
 # Similar to py_test above, this macro is used to exclude dependencies for some py_binary
 # targets in order to reduce the size of //tensorflow/tools/pip_package:simple_console_windows.
 # See https://github.com/tensorflow/tensorflow/issues/22390
@@ -2631,6 +2636,11 @@ def tf_py_test(
             **kwargs
         )
 
+register_extension_info(
+    extension = tf_py_test,
+    label_regex_for_dep = "{extension_name}(_tfrt)?",
+)
+
 def gpu_py_test(
         name,
         srcs,