Python util.assert_that方法代码示例

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_stats_pipeline_with_zero_examples(self):
    expected_result = text_format.Parse(
        """
        datasets {
          num_examples: 0
        }
        """, statistics_pb2.DatasetFeatureStatisticsList())
    with beam.Pipeline() as p:
      options = stats_options.StatsOptions(
          num_top_values=1,
          num_rank_histogram_buckets=1,
          num_values_histogram_buckets=2,
          num_histogram_buckets=1,
          num_quantiles_histogram_buckets=1,
          epsilon=0.001)
      result = (p | beam.Create([]) | stats_api.GenerateStatistics(options))
      util.assert_that(
          result,
          test_util.make_dataset_feature_stats_list_proto_equal_fn(
              self, expected_result)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_stats_pipeline_with_sample_rate(self):
    record_batches = [
        pa.RecordBatch.from_arrays(
            [pa.array([np.linspace(1, 3000, 3000, dtype=np.int32)])], ['c']),
    ]

    with beam.Pipeline() as p:
      options = stats_options.StatsOptions(
          sample_rate=1.0,
          num_top_values=2,
          num_rank_histogram_buckets=2,
          num_values_histogram_buckets=2,
          num_histogram_buckets=2,
          num_quantiles_histogram_buckets=2,
          epsilon=0.001)
      result = (
          p | beam.Create(record_batches)
          | stats_api.GenerateStatistics(options))
      util.assert_that(
          result,
          test_util.make_dataset_feature_stats_list_proto_equal_fn(
              self, self._sampling_test_expected_result)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_stats_impl(self,
                      record_batches,
                      options,
                      expected_result_proto_text,
                      schema=None):
    expected_result = text_format.Parse(
        expected_result_proto_text,
        statistics_pb2.DatasetFeatureStatisticsList())
    if schema is not None:
      options.schema = schema
    with beam.Pipeline() as p:
      result = (
          p | beam.Create(record_batches, reshuffle=False)
          | stats_impl.GenerateStatisticsImpl(options))
      util.assert_that(
          result,
          test_util.make_dataset_feature_stats_list_proto_equal_fn(
              self, expected_result)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_csv_decoder(self,
                       input_lines,
                       expected_result,
                       column_names,
                       delimiter=',',
                       skip_blank_lines=True,
                       schema=None,
                       multivalent_columns=None,
                       secondary_delimiter=None):
    with beam.Pipeline() as p:
      result = (
          p | beam.Create(input_lines, reshuffle=False)
          | csv_decoder.DecodeCSV(
              column_names=column_names,
              delimiter=delimiter,
              skip_blank_lines=skip_blank_lines,
              schema=schema,
              multivalent_columns=multivalent_columns,
              secondary_delimiter=secondary_delimiter))
      util.assert_that(
          result,
          test_util.make_arrow_record_batches_equal_fn(self, expected_result)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testReadTransformFn(self):
    path = self.get_temp_dir()
    # NOTE: we don't need to create or write to the transform_fn directory since
    # ReadTransformFn never inspects this directory.
    transform_fn_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORM_FN_DIR)
    transformed_metadata_dir = os.path.join(
        path, tft.TFTransformOutput.TRANSFORMED_METADATA_DIR)
    metadata_io.write_metadata(test_metadata.COMPLETE_METADATA,
                               transformed_metadata_dir)

    with beam.Pipeline() as pipeline:
      saved_model_dir_pcoll, metadata = (
          pipeline | transform_fn_io.ReadTransformFn(path))
      beam_test_util.assert_that(
          saved_model_dir_pcoll,
          beam_test_util.equal_to([transform_fn_dir]),
          label='AssertSavedModelDir')
      # NOTE: metadata is currently read in a non-deferred manner.
      self.assertEqual(metadata, test_metadata.COMPLETE_METADATA) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testPrestoToExample(self):
    with beam.Pipeline() as pipeline:
      examples = (
          pipeline | 'ToTFExample' >> executor._PrestoToExample(
              exec_properties={
                  'input_config':
                      json_format.MessageToJson(
                          example_gen_pb2.Input(),
                          preserving_proto_field_name=True),
                  'custom_config':
                      json_format.MessageToJson(
                          example_gen_pb2.CustomConfig(),
                          preserving_proto_field_name=True)
              },
              split_pattern='SELECT i, f, s FROM `fake`'))

      feature = {}
      feature['i'] = tf.train.Feature(int64_list=tf.train.Int64List(value=[1]))
      feature['f'] = tf.train.Feature(
          float_list=tf.train.FloatList(value=[2.0]))
      feature['s'] = tf.train.Feature(
          bytes_list=tf.train.BytesList(value=[tf.compat.as_bytes('abc')]))
      example_proto = tf.train.Example(
          features=tf.train.Features(feature=feature))
      util.assert_that(examples, util.equal_to([example_proto])) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testE2E(self, attach_raw_records):
    raw_column_name = "raw_records" if attach_raw_records else None
    tfxio = self._MakeTFXIO(_SCHEMA, raw_column_name)

    def _AssertFn(record_batch_list):
      self.assertLen(record_batch_list, 1)
      record_batch = record_batch_list[0]
      self._ValidateRecordBatch(tfxio, record_batch, raw_column_name)
      self.assertTrue(record_batch.schema.equals(tfxio.ArrowSchema()))
      tensor_adapter = tfxio.TensorAdapter()
      dict_of_tensors = tensor_adapter.ToBatchTensors(record_batch)
      self.assertLen(dict_of_tensors, 4)
      self.assertIn("int_feature", dict_of_tensors)
      self.assertIn("float_feature", dict_of_tensors)
      self.assertIn("seq_string_feature", dict_of_tensors)
      self.assertIn("seq_int_feature", dict_of_tensors)

    p = beam.Pipeline()
    record_batch_pcoll = p | tfxio.BeamSource(batch_size=1000)
    beam_testing_util.assert_that(record_batch_pcoll, _AssertFn)
    pipeline_result = p.run()
    pipeline_result.wait_until_finish()
    telemetry_test_util.ValidateMetrics(
        self, pipeline_result, _TELEMETRY_DESCRIPTORS,
        "tf_sequence_example", "tfrecords_gzip") 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testProjectionNoSequenceFeature(self):
    tfxio = self._MakeTFXIO(_SCHEMA).Project(["int_feature"])
    arrow_schema = tfxio.ArrowSchema()
    self.assertLen(arrow_schema, 1)
    self.assertIn("int_feature", arrow_schema.names)
    def _AssertFn(record_batch_list):
      self.assertLen(record_batch_list, 1)
      record_batch = record_batch_list[0]
      self._ValidateRecordBatch(tfxio, record_batch)
      tensor_adapter = tfxio.TensorAdapter()
      dict_of_tensors = tensor_adapter.ToBatchTensors(record_batch)
      self.assertLen(dict_of_tensors, 1)
      self.assertIn("int_feature", dict_of_tensors)

    with beam.Pipeline() as p:
      # Setting the betch_size to make sure only one batch is generated.
      record_batch_pcoll = p | tfxio.BeamSource(
          batch_size=len(_EXAMPLES))
      beam_testing_util.assert_that(record_batch_pcoll, _AssertFn) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_invalid_schema_type(self):
    input_lines = ['1']
    column_names = ['f1']
    schema = text_format.Parse(
        """
              feature {
                name: "struct_feature"
                type: STRUCT
              }
              """, schema_pb2.Schema())
    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
        ValueError, '.*Schema contains invalid type: STRUCT.*'):
      with beam.Pipeline() as p:
        result = (
            p | beam.Create(input_lines, reshuffle=False)
            | 'CSVToRecordBatch' >> csv_decoder.CSVToRecordBatch(
                column_names=column_names,
                schema=schema,
                desired_batch_size=1000))
        beam_test_util.assert_that(result, lambda _: None) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def test_invalid_schema_missing_column(self):
    input_lines = ['1,2']
    column_names = ['f1', 'f2']
    schema = text_format.Parse(
        """
              feature {
                name: "f1"
                type: INT
                value_count {
                  min: 0
                  max: 2
                }
              }
              """, schema_pb2.Schema())
    with self.assertRaisesRegex(  # pylint: disable=g-error-prone-assert-raises
        ValueError, '.*Schema does not contain column.*'):
      with beam.Pipeline() as p:
        result = (
            p | beam.Create(input_lines, reshuffle=False)
            | 'CSVToRecordBatch' >> csv_decoder.CSVToRecordBatch(
                column_names=column_names,
                schema=schema,
                desired_batch_size=1000))
        beam_test_util.assert_that(result, lambda _: None) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testIncludeFilter(self):
    with beam.Pipeline() as pipeline:
      got = (
          pipeline
          | 'Create' >> beam.Create([{
              'a': 1,
              'b': 2
          }])
          | 'EvaluateExtracts' >>
          analysis_table_evaluator.EvaluateExtracts(include=['a']))

      def check_result(got):
        try:
          self.assertEqual(got, [{'a': 1}])
        except AssertionError as err:
          raise util.BeamAssertException(err)

      util.assert_that(got[constants.ANALYSIS_KEY], check_result) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testExcludeFilter(self):
    with beam.Pipeline() as pipeline:
      got = (
          pipeline
          | 'Create' >> beam.Create([{
              'a': 1,
              'b': 2
          }])
          | 'EvaluateExtracts' >>
          analysis_table_evaluator.EvaluateExtracts(exclude=['a']))

      def check_result(got):
        try:
          self.assertEqual(got, [{'b': 2}])
        except AssertionError as err:
          raise util.BeamAssertException(err)

      util.assert_that(got[constants.ANALYSIS_KEY], check_result) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testNoIncludeOrExcludeFilters(self):
    with beam.Pipeline() as pipeline:
      got = (
          pipeline
          | 'Create' >> beam.Create([{
              constants.INPUT_KEY: 'input',
              'other': 2
          }])
          | 'EvaluateExtracts' >> analysis_table_evaluator.EvaluateExtracts())

      def check_result(got):
        try:
          self.assertEqual(got, [{'other': 2}])
        except AssertionError as err:
          raise util.BeamAssertException(err)

      util.assert_that(got[constants.ANALYSIS_KEY], check_result) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testIncludeFilter(self):
    with beam.Pipeline() as pipeline:
      got = (
          pipeline
          | 'Create' >> beam.Create([{
              'a': 1,
              'b': 2,
              'c': 3,
              'd': 4
          }])
          | 'Filter' >> extractor.Filter(include=['a', 'c']))

      def check_result(got):
        try:
          self.assertEqual(got, [{'a': 1, 'c': 3}])
        except AssertionError as err:
          raise util.BeamAssertException(err)

      util.assert_that(got, check_result) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import assert_that [as 别名]
def testExludeFilter(self):
    with beam.Pipeline() as pipeline:
      got = (
          pipeline
          | 'Create' >> beam.Create([{
              'a': 1,
              'b': 2,
              'c': 3,
              'd': 4
          }])
          | 'Filter' >> extractor.Filter(exclude=['b', 'd']))

      def check_result(got):
        try:
          self.assertEqual(got, [{'a': 1, 'c': 3}])
        except AssertionError as err:
          raise util.BeamAssertException(err)

      util.assert_that(got, check_result)