# Path: tensorflow_transform/beam/tft_unit.py
# def metadata_from_feature_spec(feature_spec, domains=None):
# def canonical_numeric_dtype(dtype):
# def make_feature_spec_wrapper(make_feature_spec, *args):
# def _format_example_as_numpy_dict(example, feature_shape_dict):
#     def has_ran(self):
#     def metrics(self):
#     def __exit__(self, exc_type, exc_val, exc_tb):
#   def _makeTestPipeline(self):
#   def assertMetricsCounterEqual(self, metrics, name, expected_count,
#                                 namespaces_list=None):
#   def assertAnalyzerOutputs(self,
#                             input_data,
#                             input_metadata,
#                             analyzer_fn,
#                             expected_outputs,
#                             test_data=None,
#                             desired_batch_size=None,
#                             beam_pipeline=None,
#                             force_tf_compat_v1=False,
#                             output_record_batches=False):
#     def preprocessing_fn(inputs):
#   def assertAnalyzeAndTransformResults(self,
#                                        input_data,
#                                        input_metadata,
#                                        preprocessing_fn,
#                                        expected_data=None,
#                                        expected_metadata=None,
#                                        expected_vocab_file_contents=None,
#                                        test_data=None,
#                                        desired_batch_size=None,
#                                        beam_pipeline=None,
#                                        temp_dir=None,
#                                        force_tf_compat_v1=False,
#                                        output_record_batches=False):
# class TransformTestCase(test_case.TransformTestCase):
#   class _TestPipeline(beam.Pipeline):


import os
import tensorflow_transform.beam as tft_beam
from tensorflow_transform.beam import tft_unit
# 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.
"""Tests for tensorflow_transform.beam.context."""




class ContextTest(tft_unit.TransformTestCase):

  def testNestedContextCreateBaseTempDir(self):

    level_1_dir = self.get_temp_dir()
    with tft_beam.Context(temp_dir=level_1_dir):
      self.assertEqual(
          os.path.join(level_1_dir, tft_beam.Context._TEMP_SUBDIR),
          tft_beam.Context.create_base_temp_dir())
      level_2_dir = self.get_temp_dir()
      with tft_beam.Context(temp_dir=level_2_dir):
        self.assertEqual(
            os.path.join(level_2_dir, tft_beam.Context._TEMP_SUBDIR),
            tft_beam.Context.create_base_temp_dir())
      self.assertEqual(
          os.path.join(level_1_dir, tft_beam.Context._TEMP_SUBDIR),
          tft_beam.Context.create_base_temp_dir())
    with self.assertRaises(ValueError):
      tft_beam.Context.create_base_temp_dir()


if __name__ == '__main__':

# Path: tensorflow_transform/beam/tft_unit.py
# def metadata_from_feature_spec(feature_spec, domains=None):
# def canonical_numeric_dtype(dtype):
# def make_feature_spec_wrapper(make_feature_spec, *args):
# def _format_example_as_numpy_dict(example, feature_shape_dict):
#     def has_ran(self):
#     def metrics(self):
#     def __exit__(self, exc_type, exc_val, exc_tb):
#   def _makeTestPipeline(self):
#   def assertMetricsCounterEqual(self, metrics, name, expected_count,
#                                 namespaces_list=None):
#   def assertAnalyzerOutputs(self,
#                             input_data,
#                             input_metadata,
#                             analyzer_fn,
#                             expected_outputs,
#                             test_data=None,
#                             desired_batch_size=None,
#                             beam_pipeline=None,
#                             force_tf_compat_v1=False,
#                             output_record_batches=False):
#     def preprocessing_fn(inputs):
#   def assertAnalyzeAndTransformResults(self,
#                                        input_data,
#                                        input_metadata,
#                                        preprocessing_fn,
#                                        expected_data=None,
#                                        expected_metadata=None,
#                                        expected_vocab_file_contents=None,
#                                        test_data=None,
#                                        desired_batch_size=None,
#                                        beam_pipeline=None,
#                                        temp_dir=None,
#                                        force_tf_compat_v1=False,
#                                        output_record_batches=False):
# class TransformTestCase(test_case.TransformTestCase):
#   class _TestPipeline(beam.Pipeline):


# Path: tensorflow_transform/beam/context_test.py
import os
import tensorflow_transform.beam as tft_beam
from tensorflow_transform.beam import tft_unit

# 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.
"""Tests for tensorflow_transform.beam.context."""




class ContextTest(tft_unit.TransformTestCase):

  def testNestedContextCreateBaseTempDir(self):

    level_1_dir = self.get_temp_dir()
    with tft_beam.Context(temp_dir=level_1_dir):
      self.assertEqual(
          os.path.join(level_1_dir, tft_beam.Context._TEMP_SUBDIR),
          tft_beam.Context.create_base_temp_dir())
      level_2_dir = self.get_temp_dir()
      with tft_beam.Context(temp_dir=level_2_dir):
        self.assertEqual(
            os.path.join(level_2_dir, tft_beam.Context._TEMP_SUBDIR),
            tft_beam.Context.create_base_temp_dir())
      self.assertEqual(
          os.path.join(level_1_dir, tft_beam.Context._TEMP_SUBDIR),
          tft_beam.Context.create_base_temp_dir())
    with self.assertRaises(ValueError):
      tft_beam.Context.create_base_temp_dir()


if __name__ == '__main__':

  tft_unit.main()