#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 vcfio module."""

from __future__ import absolute_import

import logging
import os
import sys
import unittest
from itertools import chain
from itertools import permutations

import apache_beam.io.source_test_utils as source_test_utils
from apache_beam.io.vcfio import _VcfSource as VcfSource
from apache_beam.io.vcfio import DEFAULT_PHASESET_VALUE
from apache_beam.io.vcfio import MISSING_GENOTYPE_VALUE
from apache_beam.io.vcfio import MalformedVcfRecord
from apache_beam.io.vcfio import ReadFromVcf
from apache_beam.io.vcfio import Variant
from apache_beam.io.vcfio import VariantCall
from apache_beam.io.vcfio import VariantInfo
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.test_utils import TempDir
from apache_beam.testing.util import BeamAssertException
from apache_beam.testing.util import assert_that

# Note: mixing \n and \r\n to verify both behaviors.
_SAMPLE_HEADER_LINES = [
    '##fileformat=VCFv4.2\n',
    '##INFO=<ID=NS,Number=1,Type=Integer,Description="Number samples">\n',
    '##INFO=<ID=AF,Number=A,Type=Float,Description="Allele Frequency">\n',
    '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\r\n',
    '##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">\n',
    '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample1	Sample2\r\n',
]

_SAMPLE_TEXT_LINES = [
    '20	14370	.	G	A	29	PASS	AF=0.5	GT:GQ	0|0:48 1|0:48\n',
    '20	17330	.	T	A	3	q10	AF=0.017	GT:GQ	0|0:49	0|1:3\n',
    '20	1110696	.	A	G,T	67	PASS	AF=0.3,0.7	GT:GQ	1|2:21	2|1:2\n',
    '20	1230237	.	T	.	47	PASS	.	GT:GQ	0|0:54	0|0:48\n',
    '19	1234567	.	GTCT	G,GTACT	50	PASS	.	GT:GQ	0/1:35	0/2:17\n',
    '20	1234	rs123	C	A,T	50	PASS	AF=0.5	GT:GQ	0/0:48	1/0:20\n',
    '19	123	rs1234	GTC	.	40	q10;s50	NS=2	GT:GQ	1|0:48	0/1:.\n',
    '19	12	.	C	<SYMBOLIC>	49	q10	AF=0.5	GT:GQ	0|1:45 .:.\n'
]


def get_full_file_path(file_name):
  """Returns the full path of the specified ``file_name`` from ``data``."""
  return os.path.join(
      os.path.dirname(__file__), '..', 'testing', 'data', 'vcf', file_name)


def get_full_dir():
  """Returns the full path of the  ``data`` directory."""
  return os.path.join(os.path.dirname(__file__), '..', 'testing', 'data', 'vcf')


# Helper method for comparing variants.
def _variant_comparator(v1, v2):
  if v1.reference_name == v2.reference_name:
    if v1.start == v2.start:
      return (v1.end > v2.end) - (v1.end < v2.end)
    return (v1.start > v2.start) - (v1.start < v2.start)
  return (v1.reference_name > v2.reference_name) - \
      (v1.reference_name < v2.reference_name)


# Helper method for verifying equal count on PCollection.
def _count_equals_to(expected_count):
  def _count_equal(actual_list):
    actual_count = len(actual_list)
    if expected_count != actual_count:
      raise BeamAssertException(
          'Expected %d not equal actual %d' % (expected_count, actual_count))
  return _count_equal


@unittest.skipIf(sys.version_info[0] == 3,
                 'VCF io will be ported to Python 3 after switch to Nucleus. '
                 'See BEAM-5628')
class VcfSourceTest(unittest.TestCase):

  # Distribution should skip tests that need VCF files due to large size
  VCF_FILE_DIR_MISSING = not os.path.exists(get_full_dir())

  def _create_temp_vcf_file(self, lines, tempdir):
    return tempdir.create_temp_file(suffix='.vcf', lines=lines)

  def _read_records(self, file_or_pattern, **kwargs):
    return source_test_utils.read_from_source(
        VcfSource(file_or_pattern, **kwargs))

  def _create_temp_file_and_read_records(self, lines):
    with TempDir() as tempdir:
      file_name = tempdir.create_temp_file(suffix='.vcf', lines=lines)
      return self._read_records(file_name)

  def _assert_variants_equal(self, actual, expected):
    self.assertEqual(
        sorted(expected),
        sorted(actual))

  def _get_sample_variant_1(self):
    """Get first sample variant.

    Features:
      multiple alternates
      not phased
      multiple names
    """
    vcf_line = ('20	1234	rs123;rs2	C	A,T	50	PASS	AF=0.5,0.1;NS=1	'
                'GT:GQ	0/0:48	1/0:20\n')
    variant = Variant(
        reference_name='20', start=1233, end=1234, reference_bases='C',
        alternate_bases=['A', 'T'], names=['rs123', 'rs2'], quality=50,
        filters=['PASS'],
        info={'AF': VariantInfo(data=[0.5, 0.1], field_count='A'),
              'NS': VariantInfo(data=1, field_count='1')})
    variant.calls.append(
        VariantCall(name='Sample1', genotype=[0, 0], info={'GQ': 48}))
    variant.calls.append(
        VariantCall(name='Sample2', genotype=[1, 0], info={'GQ': 20}))
    return variant, vcf_line

  def _get_sample_variant_2(self):
    """Get second sample variant.

    Features:
      multiple references
      no alternate
      phased
      multiple filters
      missing format field
    """
    vcf_line = (
        '19	123	rs1234	GTC	.	40	q10;s50	NS=2	GT:GQ	1|0:48	0/1:.\n')
    variant = Variant(
        reference_name='19', start=122, end=125, reference_bases='GTC',
        alternate_bases=[], names=['rs1234'], quality=40,
        filters=['q10', 's50'],
        info={'NS': VariantInfo(data=2, field_count='1')})
    variant.calls.append(
        VariantCall(name='Sample1', genotype=[1, 0],
                    phaseset=DEFAULT_PHASESET_VALUE,
                    info={'GQ': 48}))
    variant.calls.append(
        VariantCall(name='Sample2', genotype=[0, 1], info={'GQ': None}))
    return variant, vcf_line

  def _get_sample_variant_3(self):
    """Get third sample variant.

    Features:
      symbolic alternate
      no calls for sample 2
    """
    vcf_line = (
        '19	12	.	C	<SYMBOLIC>	49	q10	AF=0.5	GT:GQ	0|1:45 .:.\n')
    variant = Variant(
        reference_name='19', start=11, end=12, reference_bases='C',
        alternate_bases=['<SYMBOLIC>'], quality=49, filters=['q10'],
        info={'AF': VariantInfo(data=[0.5], field_count='A')})
    variant.calls.append(
        VariantCall(name='Sample1', genotype=[0, 1],
                    phaseset=DEFAULT_PHASESET_VALUE,
                    info={'GQ': 45}))
    variant.calls.append(
        VariantCall(name='Sample2', genotype=[MISSING_GENOTYPE_VALUE],
                    info={'GQ': None}))
    return variant, vcf_line

  def _get_invalid_file_contents(self):
    """Gets sample invalid files contents.

    Returns:
       A `tuple` where the first element is contents that are invalid because
       of record errors and the second element is contents that are invalid
       because of header errors.
    """
    malformed_vcf_records = [
        # Malfromed record.
        [
            '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample\n',
            '1    1  '
        ],
        # GT is not an integer.
        [
            '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample\n',
            '19	123	rs12345	T	C	50	q10	AF=0.2;NS=2	GT	A|0'
        ],
        # POS should be an integer.
        [
            '##FILTER=<ID=PASS,Description="All filters passed">\n',
            '##FILTER=<ID=q10,Description="Quality is less than 10.">\n',
            '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample\n',
            '19	abc	rs12345	T	C	9	q10	AF=0.2;NS=2	GT:GQ	1|0:48\n',
        ]
    ]
    malformed_header_lines = [
        # Malformed FILTER.
        [
            '##FILTER=<ID=PASS,Description="All filters passed">\n',
            '##FILTER=<ID=LowQual,Descri\n',
            '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample\n',
            '19	123	rs12345	T	C	50	q10	AF=0.2;NS=2	GT:GQ	1|0:48',
        ],
        # Invalid Number value for INFO.
        [
            '##INFO=<ID=G,Number=U,Type=String,Description="InvalidNumber">\n',
            '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO	FORMAT	Sample\n',
            '19	123	rs12345	T	C	50	q10	AF=0.2;NS=2	GT:GQ	1|0:48\n',
        ]
    ]

    return (malformed_vcf_records, malformed_header_lines)

  def test_sort_variants(self):
    sorted_variants = [
        Variant(reference_name='a', start=20, end=22),
        Variant(reference_name='a', start=20, end=22, quality=20),
        Variant(reference_name='b', start=20, end=22),
        Variant(reference_name='b', start=21, end=22),
        Variant(reference_name='b', start=21, end=23)]

    for permutation in permutations(sorted_variants):
      self.assertEqual(sorted(permutation), sorted_variants)

  def test_variant_equality(self):
    base_variant = Variant(reference_name='a', start=20, end=22,
                           reference_bases='a', alternate_bases=['g', 't'],
                           names=['variant'], quality=9, filters=['q10'],
                           info={'key': 'value'},
                           calls=[VariantCall(genotype=[0, 0])])
    equal_variant = Variant(reference_name='a', start=20, end=22,
                            reference_bases='a', alternate_bases=['g', 't'],
                            names=['variant'], quality=9, filters=['q10'],
                            info={'key': 'value'},
                            calls=[VariantCall(genotype=[0, 0])])
    different_calls = Variant(reference_name='a', start=20, end=22,
                              reference_bases='a', alternate_bases=['g', 't'],
                              names=['variant'], quality=9, filters=['q10'],
                              info={'key': 'value'},
                              calls=[VariantCall(genotype=[1, 0])])
    missing_field = Variant(reference_name='a', start=20, end=22,
                            reference_bases='a', alternate_bases=['g', 't'],
                            names=['variant'], quality=9, filters=['q10'],
                            info={'key': 'value'})

    self.assertEqual(base_variant, equal_variant)
    self.assertNotEqual(base_variant, different_calls)
    self.assertNotEqual(base_variant, missing_field)

  @unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
  def test_read_single_file_large(self):
    test_data_conifgs = [
        {'file': 'valid-4.0.vcf', 'num_records': 5},
        {'file': 'valid-4.0.vcf.gz', 'num_records': 5},
        {'file': 'valid-4.0.vcf.bz2', 'num_records': 5},
        {'file': 'valid-4.1-large.vcf', 'num_records': 9882},
        {'file': 'valid-4.2.vcf', 'num_records': 13},
    ]
    for config in test_data_conifgs:
      read_data = self._read_records(
          get_full_file_path(config['file']))
      self.assertEqual(config['num_records'], len(read_data))

  @unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
  def test_read_file_pattern_large(self):
    read_data = self._read_records(
        os.path.join(get_full_dir(), 'valid-*.vcf'))
    self.assertEqual(9900, len(read_data))
    read_data_gz = self._read_records(
        os.path.join(get_full_dir(), 'valid-*.vcf.gz'))
    self.assertEqual(9900, len(read_data_gz))

  def test_single_file_no_records(self):
    self.assertEqual(
        [], self._create_temp_file_and_read_records(['']))
    self.assertEqual(
        [], self._create_temp_file_and_read_records(['\n', '\r\n', '\n']))
    self.assertEqual(
        [], self._create_temp_file_and_read_records(_SAMPLE_HEADER_LINES))

  def test_single_file_verify_details(self):
    variant_1, vcf_line_1 = self._get_sample_variant_1()
    read_data = self._create_temp_file_and_read_records(
        _SAMPLE_HEADER_LINES + [vcf_line_1])
    self.assertEqual(1, len(read_data))
    self.assertEqual(variant_1, read_data[0])
    variant_2, vcf_line_2 = self._get_sample_variant_2()
    variant_3, vcf_line_3 = self._get_sample_variant_3()
    read_data = self._create_temp_file_and_read_records(
        _SAMPLE_HEADER_LINES + [vcf_line_1, vcf_line_2, vcf_line_3])
    self.assertEqual(3, len(read_data))
    self._assert_variants_equal([variant_1, variant_2, variant_3], read_data)

  def test_file_pattern_verify_details(self):
    variant_1, vcf_line_1 = self._get_sample_variant_1()
    variant_2, vcf_line_2 = self._get_sample_variant_2()
    variant_3, vcf_line_3 = self._get_sample_variant_3()
    with TempDir() as tempdir:
      self._create_temp_vcf_file(_SAMPLE_HEADER_LINES + [vcf_line_1], tempdir)
      self._create_temp_vcf_file((_SAMPLE_HEADER_LINES +
                                  [vcf_line_2, vcf_line_3]),
                                 tempdir)
      read_data = self._read_records(os.path.join(tempdir.get_path(), '*.vcf'))
      self.assertEqual(3, len(read_data))
      self._assert_variants_equal([variant_1, variant_2, variant_3], read_data)

  @unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
  def test_read_after_splitting(self):
    file_name = get_full_file_path('valid-4.1-large.vcf')
    source = VcfSource(file_name)
    splits = [p for p in source.split(desired_bundle_size=500)]
    self.assertGreater(len(splits), 1)
    sources_info = ([
        (split.source, split.start_position, split.stop_position) for
        split in splits])
    self.assertGreater(len(sources_info), 1)
    split_records = []
    for source_info in sources_info:
      split_records.extend(source_test_utils.read_from_source(*source_info))
    self.assertEqual(9882, len(split_records))

  def test_invalid_file(self):
    invalid_file_contents = self._get_invalid_file_contents()
    for content in chain(*invalid_file_contents):
      with TempDir() as tempdir, self.assertRaises(ValueError):
        self._read_records(self._create_temp_vcf_file(content, tempdir))
    # Try with multiple files (any one of them will throw an exception).
    with TempDir() as tempdir, self.assertRaises(ValueError):
      for content in chain(*invalid_file_contents):
        self._create_temp_vcf_file(content, tempdir)
      self._read_records(os.path.join(tempdir.get_path(), '*.vcf'))

  def test_allow_malformed_records(self):
    invalid_records, invalid_headers = self._get_invalid_file_contents()

    # Invalid records should not raise errors
    for content in invalid_records:
      with TempDir() as tempdir:
        records = self._read_records(
            self._create_temp_vcf_file(content, tempdir),
            allow_malformed_records=True)
        for record in records:
          self.assertIsInstance(record, MalformedVcfRecord)

    # Invalid headers should still raise errors
    for content in invalid_headers:
      with TempDir() as tempdir, self.assertRaises(ValueError):
        self._read_records(self._create_temp_vcf_file(content, tempdir),
                           allow_malformed_records=True)

  def test_no_samples(self):
    header_line = '#CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO\n'
    record_line = '19	123	.	G	A	.	PASS	AF=0.2'
    expected_variant = Variant(
        reference_name='19', start=122, end=123, reference_bases='G',
        alternate_bases=['A'], filters=['PASS'],
        info={'AF': VariantInfo(data=[0.2], field_count='A')})
    read_data = self._create_temp_file_and_read_records(
        _SAMPLE_HEADER_LINES[:-1] + [header_line, record_line])
    self.assertEqual(1, len(read_data))
    self.assertEqual(expected_variant, read_data[0])

  def test_no_info(self):
    record_line = 'chr19	123	.	.	.	.	.	.	GT	.	.'
    expected_variant = Variant(reference_name='chr19', start=122, end=123)
    expected_variant.calls.append(
        VariantCall(name='Sample1', genotype=[MISSING_GENOTYPE_VALUE]))
    expected_variant.calls.append(
        VariantCall(name='Sample2', genotype=[MISSING_GENOTYPE_VALUE]))
    read_data = self._create_temp_file_and_read_records(
        _SAMPLE_HEADER_LINES + [record_line])
    self.assertEqual(1, len(read_data))
    self.assertEqual(expected_variant, read_data[0])

  def test_info_numbers_and_types(self):
    info_headers = [
        '##INFO=<ID=HA,Number=A,Type=String,Description="StringInfo_A">\n',
        '##INFO=<ID=HG,Number=G,Type=Integer,Description="IntInfo_G">\n',
        '##INFO=<ID=HR,Number=R,Type=Character,Description="ChrInfo_R">\n',
        '##INFO=<ID=HF,Number=0,Type=Flag,Description="FlagInfo">\n',
        '##INFO=<ID=HU,Number=.,Type=Float,Description="FloatInfo_variable">\n']
    record_lines = [
        '19	2	.	A	T,C	.	.	HA=a1,a2;HG=1,2,3;HR=a,b,c;HF;HU=0.1	GT	1/0	0/1\n',
        '19	124	.	A	T	.	.	HG=3,4,5;HR=d,e;HU=1.1,1.2	GT	0/0	0/1']
    variant_1 = Variant(
        reference_name='19', start=1, end=2, reference_bases='A',
        alternate_bases=['T', 'C'],
        info={'HA': VariantInfo(data=['a1', 'a2'], field_count='A'),
              'HG': VariantInfo(data=[1, 2, 3], field_count='G'),
              'HR': VariantInfo(data=['a', 'b', 'c'], field_count='R'),
              'HF': VariantInfo(data=True, field_count='0'),
              'HU': VariantInfo(data=[0.1], field_count=None)})
    variant_1.calls.append(VariantCall(name='Sample1', genotype=[1, 0]))
    variant_1.calls.append(VariantCall(name='Sample2', genotype=[0, 1]))
    variant_2 = Variant(
        reference_name='19', start=123, end=124, reference_bases='A',
        alternate_bases=['T'],
        info={'HG': VariantInfo(data=[3, 4, 5], field_count='G'),
              'HR': VariantInfo(data=['d', 'e'], field_count='R'),
              'HU': VariantInfo(data=[1.1, 1.2], field_count=None)})
    variant_2.calls.append(VariantCall(name='Sample1', genotype=[0, 0]))
    variant_2.calls.append(VariantCall(name='Sample2', genotype=[0, 1]))
    read_data = self._create_temp_file_and_read_records(
        info_headers + _SAMPLE_HEADER_LINES[1:] + record_lines)
    self.assertEqual(2, len(read_data))
    self._assert_variants_equal([variant_1, variant_2], read_data)

  def test_end_info_key(self):
    phaseset_header_line = (
        '##INFO=<ID=END,Number=1,Type=Integer,Description="End of record.">\n')
    record_lines = ['19	123	.	A	.	.	.	END=1111	GT	1/0	0/1\n',
                    '19	123	.	A	.	.	.	.	GT	0/1	1/1\n']
    variant_1 = Variant(
        reference_name='19', start=122, end=1111, reference_bases='A')
    variant_1.calls.append(VariantCall(name='Sample1', genotype=[1, 0]))
    variant_1.calls.append(VariantCall(name='Sample2', genotype=[0, 1]))
    variant_2 = Variant(
        reference_name='19', start=122, end=123, reference_bases='A')
    variant_2.calls.append(VariantCall(name='Sample1', genotype=[0, 1]))
    variant_2.calls.append(VariantCall(name='Sample2', genotype=[1, 1]))
    read_data = self._create_temp_file_and_read_records(
        [phaseset_header_line] + _SAMPLE_HEADER_LINES[1:] + record_lines)
    self.assertEqual(2, len(read_data))
    self._assert_variants_equal([variant_1, variant_2], read_data)

  def test_custom_phaseset(self):
    phaseset_header_line = (
        '##FORMAT=<ID=PS,Number=1,Type=Integer,Description="Phaseset">\n')
    record_lines = ['19	123	.	A	T	.	.	.	GT:PS	1|0:1111	0/1:.\n',
                    '19	121	.	A	T	.	.	.	GT:PS	1|0:2222	0/1:2222\n']
    variant_1 = Variant(
        reference_name='19', start=122, end=123, reference_bases='A',
        alternate_bases=['T'])
    variant_1.calls.append(
        VariantCall(name='Sample1', genotype=[1, 0], phaseset='1111'))
    variant_1.calls.append(VariantCall(name='Sample2', genotype=[0, 1]))
    variant_2 = Variant(
        reference_name='19', start=120, end=121, reference_bases='A',
        alternate_bases=['T'])
    variant_2.calls.append(
        VariantCall(name='Sample1', genotype=[1, 0], phaseset='2222'))
    variant_2.calls.append(
        VariantCall(name='Sample2', genotype=[0, 1], phaseset='2222'))
    read_data = self._create_temp_file_and_read_records(
        [phaseset_header_line] + _SAMPLE_HEADER_LINES[1:] + record_lines)
    self.assertEqual(2, len(read_data))
    self._assert_variants_equal([variant_1, variant_2], read_data)

  def test_format_numbers(self):
    format_headers = [
        '##FORMAT=<ID=FU,Number=.,Type=String,Description="Format_variable">\n',
        '##FORMAT=<ID=F1,Number=1,Type=Integer,Description="Format_one">\n',
        '##FORMAT=<ID=F2,Number=2,Type=Character,Description="Format_two">\n']
    record_lines = [
        '19	2	.	A	T,C	.	.	.	GT:FU:F1:F2	1/0:a1:3:a,b	0/1:a2,a3:4:b,c\n']
    expected_variant = Variant(
        reference_name='19', start=1, end=2, reference_bases='A',
        alternate_bases=['T', 'C'])
    expected_variant.calls.append(VariantCall(
        name='Sample1',
        genotype=[1, 0],
        info={'FU': ['a1'], 'F1': 3, 'F2': ['a', 'b']}))
    expected_variant.calls.append(VariantCall(
        name='Sample2',
        genotype=[0, 1],
        info={'FU': ['a2', 'a3'], 'F1': 4, 'F2': ['b', 'c']}))
    read_data = self._create_temp_file_and_read_records(
        format_headers + _SAMPLE_HEADER_LINES[1:] + record_lines)
    self.assertEqual(1, len(read_data))
    self.assertEqual(expected_variant, read_data[0])

  def test_pipeline_read_single_file(self):
    with TempDir() as tempdir:
      file_name = self._create_temp_vcf_file(_SAMPLE_HEADER_LINES +
                                             _SAMPLE_TEXT_LINES, tempdir)
      pipeline = TestPipeline()
      pcoll = pipeline | 'Read' >> ReadFromVcf(file_name)
      assert_that(pcoll, _count_equals_to(len(_SAMPLE_TEXT_LINES)))
      pipeline.run()

  @unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
  def test_pipeline_read_single_file_large(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Read' >> ReadFromVcf(
        get_full_file_path('valid-4.0.vcf'))
    assert_that(pcoll, _count_equals_to(5))
    pipeline.run()

  @unittest.skipIf(VCF_FILE_DIR_MISSING, 'VCF test file directory is missing')
  def test_pipeline_read_file_pattern_large(self):
    pipeline = TestPipeline()
    pcoll = pipeline | 'Read' >> ReadFromVcf(
        os.path.join(get_full_dir(), 'valid-*.vcf'))
    assert_that(pcoll, _count_equals_to(9900))
    pipeline.run()

  def test_read_reentrant_without_splitting(self):
    with TempDir() as tempdir:
      file_name = self._create_temp_vcf_file(_SAMPLE_HEADER_LINES +
                                             _SAMPLE_TEXT_LINES, tempdir)
      source = VcfSource(file_name)
      source_test_utils.assert_reentrant_reads_succeed((source, None, None))

  def test_read_reentrant_after_splitting(self):
    with TempDir() as tempdir:
      file_name = self._create_temp_vcf_file(_SAMPLE_HEADER_LINES +
                                             _SAMPLE_TEXT_LINES, tempdir)
      source = VcfSource(file_name)
      splits = [split for split in source.split(desired_bundle_size=100000)]
      assert len(splits) == 1
      source_test_utils.assert_reentrant_reads_succeed(
          (splits[0].source, splits[0].start_position, splits[0].stop_position))

  def test_dynamic_work_rebalancing(self):
    with TempDir() as tempdir:
      file_name = self._create_temp_vcf_file(_SAMPLE_HEADER_LINES +
                                             _SAMPLE_TEXT_LINES, tempdir)
      source = VcfSource(file_name)
      splits = [split for split in source.split(desired_bundle_size=100000)]
      assert len(splits) == 1
      source_test_utils.assert_split_at_fraction_exhaustive(
          splits[0].source, splits[0].start_position, splits[0].stop_position)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()