# Path: tensorflow_transform/common.py
# ANALYZER_COLLECTION = 'tft_analyzer_use'
# MAPPER_COLLECTION = 'tft_mapper_use'
# ANNOTATION_PREFIX_URL = 'type.googleapis.com'
#   IS_ANNOTATIONS_PB_AVAILABLE = True
#   IS_ANNOTATIONS_PB_AVAILABLE = False
# def logging_context() -> Generator[None, None, None]:
# def log_api_use(
#     collection_name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
#   def decorator(fn):
#     def wrapped_fn(*args, **kwargs):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


import tensorflow as tf
from tensorflow_transform import common
from tensorflow_transform import test_case

    self.assertAllEqual([{"fn0": 3, "fn1": 1}],
                        graph.get_collection("test_collection"))
    self.assertAllEqual([{"fn2": 1}],
                        graph.get_collection("another_collection"))

  def testLogAPIUseWithNestedFunction(self):
    """Tests that API call is not logged when called from another logged API."""

    @common.log_api_use("test_collection")
    def fn0():
      fn1()
      return fn2()

    @common.log_api_use("test_collection")
    def fn1():
      return None

    @common.log_api_use("another_collection")
    def fn2():
      return None

    with tf.compat.v1.Graph().as_default() as graph:
      fn0()

    self.assertEqual([{"fn0": 1}], graph.get_collection("test_collection"))
    self.assertAllEqual([], graph.get_collection("another_collection"))


if __name__ == "__main__":

# Path: tensorflow_transform/common.py
# ANALYZER_COLLECTION = 'tft_analyzer_use'
# MAPPER_COLLECTION = 'tft_mapper_use'
# ANNOTATION_PREFIX_URL = 'type.googleapis.com'
#   IS_ANNOTATIONS_PB_AVAILABLE = True
#   IS_ANNOTATIONS_PB_AVAILABLE = False
# def logging_context() -> Generator[None, None, None]:
# def log_api_use(
#     collection_name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
#   def decorator(fn):
#     def wrapped_fn(*args, **kwargs):
# 
# Path: tensorflow_transform/test_case.py
# def is_tf_api_version_1():
# def cross_named_parameters(*args):
#   def _cross_test_cases(parameters_list):
# def parameters(*testcases):
#   def wrapper(fn):
#     def to_arg_dict(testcase):
# def cross_parameters(*args):
# def _make_placeholder(tensor_spec):
# def _graph_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_graph(*inputs):
# def _ragged_value_as_constant(value, dtype):
# def _wrap_as_constant(value, tensor_spec):
# def _eager_function_handler(input_signature):
#   def wrapper(fn):
#     def _run_eagerly(*inputs):  # pylint: disable=missing-docstring
# def _tf_function_function_handler(input_signature):
#   def wrapper(fn):
# def is_external_environment():
# def skip_if_external_environment(reason):
# def skip_if_not_tf2(reason):
# def cross_with_function_handlers(parameters_list):
#   def assertDataCloseOrEqual(self, a_data, b_data):
#   def _assertValuesCloseOrEqual(self, a_value, b_value, msg=None):
#   def AssertVocabularyContents(self, vocab_file_path, file_contents):
#   def WriteRenderedDotFile(self, dot_string, output_file=None):
#   def _NumpyArraysToLists(self, maybe_arrays):
#   def _SortedDicts(self, list_of_dicts):
#   def _SortedData(self, list_of_dicts_of_arrays):
# FUNCTION_HANDLERS = [
#     dict(testcase_name='graph',
#          function_handler=_graph_function_handler),
#     dict(testcase_name='eager',
#          function_handler=_eager_function_handler),
#     dict(testcase_name='tf_function',
#          function_handler=_tf_function_function_handler)
# ]
# class TransformTestCase(parameterized.TestCase, tf.test.TestCase):


# Path: tensorflow_transform/common_test.py
import tensorflow as tf
from tensorflow_transform import common
from tensorflow_transform import test_case


    self.assertAllEqual([{"fn0": 3, "fn1": 1}],
                        graph.get_collection("test_collection"))
    self.assertAllEqual([{"fn2": 1}],
                        graph.get_collection("another_collection"))

  def testLogAPIUseWithNestedFunction(self):
    """Tests that API call is not logged when called from another logged API."""

    @common.log_api_use("test_collection")
    def fn0():
      fn1()
      return fn2()

    @common.log_api_use("test_collection")
    def fn1():
      return None

    @common.log_api_use("another_collection")
    def fn2():
      return None

    with tf.compat.v1.Graph().as_default() as graph:
      fn0()

    self.assertEqual([{"fn0": 1}], graph.get_collection("test_collection"))
    self.assertAllEqual([], graph.get_collection("another_collection"))


if __name__ == "__main__":

  test_case.main()