Python util.equal_to方法代码示例

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_shard_variants(self):
    expected_shards = self._get_expected_variant_shards()
    variants = [variant
                for variant_list in expected_shards.values()
                for variant in variant_list]

    sharding = variant_sharding.VariantSharding(
        'gcp_variant_transforms/data/sharding_configs/'
        'homo_sapiens_default.yaml')
    pipeline = TestPipeline()
    shards = (
        pipeline
        | Create(variants, reshuffle=False)
        | 'ShardVariants' >> beam.Partition(
            shard_variants.ShardVariants(sharding),
            sharding.get_num_shards()))
    for i in range(sharding.get_num_shards()):
      assert_that(shards[i], equal_to(expected_shards.get(i, [])),
                  label=str(i))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_merge_header_definitions_one_header(self):
    lines = [
        '##INFO=<ID=NS,Number=1,Type=Integer,Description="Number samples">\n',
        '#CHROM  POS ID  REF ALT QUAL  FILTER  INFO  FORMAT  Sample1 Sample2\n'
    ]

    headers = self._get_header_from_lines(lines, 'file1')
    pipeline = TestPipeline()
    merged_definitions = (
        pipeline
        | Create([headers])
        | 'MergeDefinitions' >> merge_header_definitions.MergeDefinitions())

    expected = VcfHeaderDefinitions()
    expected._infos = {'NS': {Definition(1, 'Integer'): ['file1']}}
    assert_that(merged_definitions, equal_to([expected]))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_merge_header_definitions_two_conflicting_headers(self):
    lines_1 = [
        '##INFO=<ID=NS,Number=1,Type=Integer,Description="Number samples">\n',
        '#CHROM  POS ID  REF ALT QUAL  FILTER  INFO  FORMAT  Sample1 Sample2\n'
    ]
    lines_2 = [
        '##INFO=<ID=NS,Number=1,Type=Float,Description="Number samples">\n',
        '#CHROM  POS ID  REF ALT QUAL  FILTER  INFO  FORMAT  Sample3\n'
    ]

    headers_1 = self._get_header_from_lines(lines_1, 'file1')
    headers_2 = self._get_header_from_lines(lines_2, 'file2')
    pipeline = TestPipeline()
    merged_definitions = (
        pipeline
        | Create([headers_1, headers_2])
        | 'MergeDefinitions' >> merge_header_definitions.MergeDefinitions())

    expected = VcfHeaderDefinitions()
    expected._infos = {'NS': {Definition(1, 'Integer'): ['file1'],
                              Definition(1, 'Float'): ['file2']}}
    assert_that(merged_definitions, equal_to([expected]))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_merge_header_definitions_same_id_in_info_and_format_headers(self):
    lines_1 = [
        '##INFO=<ID=NS,Number=1,Type=Integer,Description="Number samples">\n',
        '#CHROM  POS ID  REF ALT QUAL  FILTER  INFO  FORMAT  Sample1 Sample2\n'
    ]
    lines_2 = [
        '##FORMAT=<ID=NS,Number=1,Type=Float,Description="Number samples">\n',
        '#CHROM  POS ID  REF ALT QUAL  FILTER  INFO  FORMAT  Sample3\n'
    ]

    headers_1 = self._get_header_from_lines(lines_1, 'file1')
    headers_2 = self._get_header_from_lines(lines_2, 'file2')
    pipeline = TestPipeline()
    merged_definitions = (
        pipeline
        | Create([headers_1, headers_2])
        | 'MergeDefinitions' >> merge_header_definitions.MergeDefinitions())

    expected = VcfHeaderDefinitions()
    expected._infos = {'NS': {Definition(1, 'Integer'): ['file1']}}
    expected._formats = {'NS': {Definition(1, 'Float'): ['file2']}}

    assert_that(merged_definitions, equal_to([expected]))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_convert_variant_to_bigquery_row(self):
    variant_1, row_1, header_num_dict_1 = self._get_sample_variant_1()
    variant_2, row_2, header_num_dict_2 = self._get_sample_variant_2()
    variant_3, row_3, header_num_dict_3 = self._get_sample_variant_3()
    header_num_dict = header_num_dict_1.copy()
    header_num_dict.update(header_num_dict_2)
    header_num_dict.update(header_num_dict_3)
    header_fields = vcf_header_util.make_header(header_num_dict)
    proc_var_1 = processed_variant.ProcessedVariantFactory(
        header_fields).create_processed_variant(variant_1)
    proc_var_2 = processed_variant.ProcessedVariantFactory(
        header_fields).create_processed_variant(variant_2)
    proc_var_3 = processed_variant.ProcessedVariantFactory(
        header_fields).create_processed_variant(variant_3)
    pipeline = TestPipeline(blocking=True)
    bigquery_rows = (
        pipeline
        | Create([proc_var_1, proc_var_2, proc_var_3])
        | 'ConvertToRow' >> beam.ParDo(ConvertVariantToRow(
            self._row_generator)))
    assert_that(bigquery_rows, equal_to([row_1, row_2, row_3]))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_sample_ids_combiner_pipeline(self):
    sample_ids = [hash_name('sample3'),
                  hash_name('sample2'),
                  hash_name('sample1')]
    variant_calls = [
        vcfio.VariantCall(sample_id=sample_ids[0]),
        vcfio.VariantCall(sample_id=sample_ids[1]),
        vcfio.VariantCall(sample_id=sample_ids[2])
    ]
    variants = [
        vcfio.Variant(calls=[variant_calls[0], variant_calls[1]]),
        vcfio.Variant(calls=[variant_calls[1], variant_calls[2]])
    ]

    pipeline = TestPipeline()
    combined_sample_ids = (
        pipeline
        | transforms.Create(variants)
        | 'CombineSampleIds' >> combine_sample_ids.SampleIdsCombiner()
        | combiners.ToList())
    assert_that(combined_sample_ids, equal_to([sample_ids]))
    pipeline.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_header_fields_inferred_one_variant(self):
    with TestPipeline() as p:
      variant = self._get_sample_variant_1()
      inferred_headers = (
          p
          | Create([variant])
          | 'InferHeaderFields' >>
          infer_headers.InferHeaderFields(defined_headers=None,
                                          infer_headers=True))

      expected_infos = {'IS': createInfo('IS', 1, 'String', ''),
                        'ISI': createInfo('ISI', 1, 'Integer', ''),
                        'ISF': createInfo('ISF', 1, 'Float', ''),
                        'IF': createInfo('IF', 1, 'Float', ''),
                        'IB': createInfo('IB', 0, 'Flag', ''),
                        'IA': createInfo('IA', '.', 'Integer', '')}
      expected_formats = {'FI': createFormat('FI', 1, 'Integer', ''),
                          'FU': createFormat('FU', '.', 'Float', '')}

      expected = vcf_header_io.VcfHeader(
          infos=expected_infos, formats=expected_formats)
      assert_that(inferred_headers, equal_to([expected]))
      p.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_defined_fields_filtered_one_variant(self):
    # All FORMATs and INFOs are already defined in the header section of VCF
    # files.
    with TestPipeline() as p:
      vcf_headers = self._get_sample_header_fields()
      vcf_headers_side_input = p | 'vcf_headers' >> Create([vcf_headers])
      variant = self._get_sample_variant_1()
      inferred_headers = (
          p
          | Create([variant])
          | 'InferHeaderFields' >>
          infer_headers.InferHeaderFields(
              pvalue.AsSingleton(vcf_headers_side_input),
              infer_headers=True))
      expected = vcf_header_io.VcfHeader()
      assert_that(inferred_headers, equal_to([expected]))
      p.run() 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [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 equal_to [as 别名]
def testTwoLangs(self):
    with TestPipeline() as p:
      tokens = p | 'CreateInput' >> beam.Create(self.sample_input)
      result = tokens | beam.ParDo(utils.CompileTokenizationInfo())
      assert_that(result, equal_to([{
          'lang': 'en',
          'count': 1,
          'num_preserved_chars': 13,
          'num_dropped_chars': 2,
          'num_non_unk_wordpieces': 4,
          'preserved_ratio': [13/4],
          'dropped_ratio': [2/15],
          'wordpieces': collections.Counter(['the', 'app', '##le', 'sauce'])
      }, {
          'lang': 'fr',
          'count': 1,
          'num_preserved_chars': 14,
          'num_dropped_chars': 0,
          'num_non_unk_wordpieces': 5,
          'preserved_ratio': [14/5],
          'dropped_ratio': [0],
          'wordpieces': collections.Counter(['bon', '##jour', 'bon', '##soir'])
      }])) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [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 equal_to [as 别名]
def test_handle_null_user_tags(self):
    """Tests the _handle_null_user_tags function."""
    source = [
        {constants.USER_TAGS_KEY: []},
        {constants.USER_TAGS_KEY: [1] * constants.USER_TAGS_LENGTH},
    ]
    target = [
        {constants.USER_TAGS_KEY: [0] * constants.USER_TAGS_LENGTH},
        {constants.USER_TAGS_KEY: [1] * constants.USER_TAGS_LENGTH},
    ]

    with TestPipeline() as p:
      features = (p
                  | "CreateStubs" >> beam.Create(source)
                  | "HandleNullUserTags" >> beam.Map(
                      preprocess._handle_null_user_tags))
      assert_that(features, equal_to(target)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_normalize_user_tags(self):
    """Tests the _normalize_user_tags function."""
    n_tags = constants.USER_TAGS_LENGTH
    source = [
        {constants.USER_TAGS_KEY: [0] * n_tags},
        {constants.USER_TAGS_KEY: [1] * n_tags},
    ]
    target = [
        {constants.USER_TAGS_KEY: [0] * n_tags},
        {constants.USER_TAGS_KEY: [n_tags**-1] * n_tags},
    ]

    with TestPipeline() as p:
      features = (p
                  | "CreateStubs" >> beam.Create(source)
                  | "NormalizeUserTags" >> beam.Map(
                      preprocess._normalize_user_tags))
      assert_that(features, equal_to(target)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_append_lifetime_duration(self):
        """Test _append_life_duration step.

        Expected behavior is that the pipeline appends the calculated lifetime
        duration (in days) of each user to a list of dicts containing BigQuery
        output.
        """
        with TestPipeline() as p:
            input_data = (
                p
                | 'CreateSourceData' >> beam.Create(
                    PreprocessTest.SOURCE_DATA))

            duration_data = (
                input_data
                | 'AppendLifetime' >> beam.Map(
                    preprocess.append_lifetime_duration))

            assert_that(duration_data, equal_to(
                PreprocessTest.DURATION_DATA)) 

# 需要导入模块: from apache_beam.testing import util [as 别名]
# 或者: from apache_beam.testing.util import equal_to [as 别名]
def test_append_label(self):
        """Test _append_label step.

        Expected behavior is that the pipeline appends the String label (i.e.
        '0-2M') of each user to a list of dicts containing BigQuery output.
        """
        with TestPipeline() as p:
            input_data = (
                p
                | 'CreateDurationData' >> beam.Create(
                    PreprocessTest.DURATION_DATA))

            label_data = (
                input_data
                | 'AppendLabel' >> beam.Map(
                    preprocess.append_label))

            assert_that(label_data, equal_to(
                PreprocessTest.LABEL_DATA))