Merge pull request #176 from ncsa/172-fix-vcf-headers

Fixed and cleaned up tests

Author: joshfactorial

neat/common/io.py

@@ -63,10 +63,9 @@ def open_input(path: str | Path) -> Iterator[TextIO]:
     # - https://github.com/python/mypy/issues/12053
     open_: Callable[..., TextIO]
     if is_compressed(path):
-        open_ = bgzf.open
+        handle = bgzf.BgzfReader(path, 'r')
     else:
-        open_ = open
-    handle = open_(path, "rt", encoding="utf-8")
+        handle = open(path, 'rt', encoding='utf-8')
     try:
         yield handle
     finally:

neat/read_simulator/utils/__init__.py

@@ -5,7 +5,6 @@
 from .options import *
 from .output_file_writer import *
 from .read import *
-from .t_sam_record import *
 from .vcf_func import *
 from .generate_reads import *
 from .generate_variants import *

neat/read_simulator/utils/stitch_outputs.py

@@ -21,28 +21,27 @@
 __all__ = ["main"]
 
 from Bio import SeqIO, bgzf
+from Bio.bgzf import BgzfWriter
 
 from neat.common import open_output, open_input
 from neat.read_simulator.utils import Options, OutputFileWriter
 
 _LOG = logging.getLogger(__name__)
 
-def concat(files_to_join: List[Path], ofw: OutputFileWriter, file: Path) -> None:
+def concat(files_to_join: List[Path], dest_file: BgzfWriter) -> None:
     if not files_to_join:
         # Nothing to do, and no error to throw
         return
 
-    out_handle = ofw.files_to_write[file]
     for f in files_to_join:
         with bgzf.BgzfReader(f) as in_f:
-            shutil.copyfileobj(in_f, out_handle)
+            shutil.copyfileobj(in_f, dest_file)
 
 def merge_bam(bam_files: List[Path], ofw: OutputFileWriter, threads: int) -> None:
     if not bam_files:
         return
 
     unsorted = ofw.bam.with_suffix(".unsorted.bam")
-    sorted_bam_files = []
     pysam.merge("--no-PG", "-@", str(threads), "-f", str(unsorted), *map(str, bam_files))
     pysam.sort("-@", str(threads), "-o", str(ofw.bam), str(unsorted))
     unsorted.unlink(missing_ok=True)
@@ -65,8 +64,8 @@ def main(
         if file_dict["bam"]:
             bam.append(file_dict["bam"])
     # concatenate all files of each type. An empty list will result in no action
-    concat(fq1_list, ofw, ofw.fq1)
-    concat(fq2_list, ofw, ofw.fq2)
+    concat(fq1_list, ofw.files_to_write[ofw.fq1])
+    concat(fq2_list, ofw.files_to_write[ofw.fq2])
     merge_bam(bam, ofw, threads)
     # Final success message via logging
     _LOG.info("Stitching complete!")

neat/read_simulator/utils/t_sam_record.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import pytest
+from Bio import bgzf
 
 from neat.common.io import is_compressed, open_input, open_output, validate_input_path, validate_output_path
 
@@ -12,7 +13,7 @@ def test_is_compressed_plain_and_gz(tmp_path: Path):
     plain = tmp_path / "file.txt"
     plain.write_text("hello", encoding="utf-8")
     gz = tmp_path / "file.txt.gz"
-    with gzip.open(gz, "wt", encoding="utf-8") as f:
+    with bgzf.BgzfWriter(gz, "wt") as f:
         f.write("hello")
 
     assert is_compressed(plain) is False
@@ -28,10 +29,13 @@ def test_open_input_reads_plain_and_gz(tmp_path: Path):
 
     # Gz
     gz = tmp_path / "b.txt.gz"
-    with gzip.open(gz, "wt", encoding="utf-8") as fh:
+    with bgzf.BgzfWriter(gz, "wt") as fh:
         fh.write("xyz\n")
+    out = ""
     with open_input(gz) as fh:
-        assert fh.read() == "xyz\n"
+        for line in fh:
+            out += line
+    assert out == "xyz\n"
 
 
 def test_open_output_creates_dirs_and_writes_plain(tmp_path: Path):

tests/test_common/test_io.py

@@ -1,19 +1,21 @@
 """
 Unit tests for the split_inputs module of the parallel read simulator
 """
+from Bio.Seq import Seq
 
-from neat.read_simulator.utils.split_inputs import SimpleRecord, chunk_record
-
+from neat.read_simulator.utils.split_inputs import chunk_record
 
 def test_chunk_record_overlaps() -> None:
     """Ensure that chunk_record yields overlapping chunks with the correct ids."""
     # A simple sequence of 12 bases to chunk into length 5 with overlap 2
-    rec = SimpleRecord("contig", "ACGTACGTACGT")
+    rec = Seq("ACGTACGTACGT")
     chunks = list(chunk_record(rec, 5, 2))
     # Should produce four chunks: positions [0:5], [3:8], [6:11], [9:12]
     assert len(chunks) == 4
     # Check that chunk ids are sequential and lengths match expected slices
-    lengths = [len(r.seq) for r, _ in chunks]
+    lengths = [len(seq) for _, seq in chunks]
     assert lengths == [5, 5, 5, 3]
-    ids = [cid for _, cid in chunks]
-    assert ids == [1, 2, 3, 4]
+    ids = [cid for cid, _ in chunks]
+    seqs = [seq for _, seq in chunks]
+    assert ids == [0, 3, 6, 9]
+    assert seqs == [Seq('ACGTA'), Seq('TACGT'), Seq('GTACG'), Seq('CGT')]

tests/test_models/test_parallelize.py

@@ -3,24 +3,35 @@
 """
 
 from pathlib import Path
+
+from Bio import bgzf
+
 from neat.read_simulator.utils.stitch_outputs import concat
 
 
 def test_concat_joins_files_in_order(tmp_path: Path) -> None:
     """Verify that concat writes the exact bytewise concatenation of inputs."""
     # Prepare small input files
     f1 = tmp_path / "a.bin"
-    f1.write_bytes(b"hello\n")
+    with bgzf.BgzfWriter(f1, 'w')as f1_in:
+        f1_in.write("hello\n")
     f2 = tmp_path / "b.bin"
-    f2.write_bytes(b"world\n")
+    with bgzf.BgzfWriter(f2, 'w') as f2_in:
+        f2_in.write("world\n")
     f3 = tmp_path / "c.bin"
-    f3.write_bytes(b"!!!")
+    with bgzf.BgzfWriter(f3, 'w') as f3_in:
+        f3_in.write("!!!")
 
     dest = tmp_path / "out.bin"
-    concat([f1, f2, f3], dest)
-
+    dest_write = bgzf.BgzfWriter(dest, 'w')
+    concat([f1, f2, f3], dest_write)
+    dest_write.close()
     assert dest.exists()
-    assert dest.read_bytes() == b"hello\nworld\n!!!"
+    with bgzf.BgzfReader(dest) as read_dest:
+        text = ""
+        for line in read_dest:
+            text += line
+    assert text == "hello\nworld\n!!!"
 
 
 def test_concat_noop_on_empty_list(tmp_path: Path) -> None:

tests/test_models/test_split_inputs.py

@@ -83,27 +83,25 @@ def test_various_read_lengths():
 
 def test_fragment_mean_st_dev_combinations():
     """Test cover_dataset with combinations of fragment mean and standard deviation to ensure no errors"""
-    span_length = 10000
+    span_length = 5000
     options = Options(rng_seed=0)
-    options.paired_ended = True
-    options.read_len = 100
-    options.coverage = 5
+    options.paired_ended = False
+    options.read_len = 101
+    options.coverage = 2
     options.overwrite_output = True
 
-    fragment_means = [100, 150, 200, 250, 300, 350, 400, 450, 500, 750, 1000]
-    fragment_st_devs = [1, 2, 5, 10, 25, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 750, 1000]
+    fragment_means = [100, 150, 200, 250,]
+    fragment_st_devs = [1, 5, 25, 50]
 
     for mean in fragment_means:
         for st_dev in fragment_st_devs:
             options.fragment_mean = mean
             options.fragment_st_dev = st_dev
             fragment_model = FragmentLengthModel(mean, st_dev)
-            try:
-                reads = cover_dataset(span_length, options, fragment_model)
-                assert isinstance(reads, list)
-            except Exception as e:
-                pytest.fail(f"Test failed for mean={mean}, st_dev={st_dev} with exception: {e}")
-
+            frags = fragment_model.generate_fragments(20, options.rng)
+            assert len(frags) == 20
+            assert fragment_model.fragment_mean == mean
+            assert fragment_model.fragment_st_dev == st_dev
 
 def test_coverage_ploidy_combinations():
     """Test cover_dataset with various combinations of coverage and ploidy values to ensure no errors"""
@@ -116,7 +114,7 @@ def test_coverage_ploidy_combinations():
     options.overwrite_output = True
     fragment_model = FragmentLengthModel(250, 100)
 
-    coverage_values = [1, 2, 5, 10, 25, 50]
+    coverage_values = [1, 2, 5, 10]
     ploidy_values = [1, 2, 4]
 
     for coverage in coverage_values: