Fixup tests

Author: jeromekelleher

bio2zarr/vcf2zarr/icf.py

@@ -354,13 +354,13 @@ def scan_vcfs(paths, show_progress, target_num_partitions, worker_processes=1):
     # Note: this will be infinity here if any of the chunks has an index
     # that doesn't keep track of the number of records per-contig
     icf_metadata.num_records = total_records
-    if len(all_partitions) > 1:
-        # Sort by contig (in the order they appear in the header) first,
-        # then by start coordinate
-        contig_index_map = {contig.id: j for j, contig in enumerate(metadata.contigs)}
-        all_partitions.sort(
-            key=lambda x: (contig_index_map[x.region.contig], x.region.start)
-        )
+
+    # Sort by contig (in the order they appear in the header) first,
+    # then by start coordinate
+    contig_index_map = {contig.id: j for j, contig in enumerate(metadata.contigs)}
+    all_partitions.sort(
+        key=lambda x: (contig_index_map[x.region.contig], x.region.start)
+    )
     icf_metadata.partitions = all_partitions
     logger.info(f"Scan complete, resulting in {len(all_partitions)} partitions.")
     return icf_metadata, header

bio2zarr/vcf_utils.py

@@ -90,7 +90,9 @@ class Region:
 
     def __post_init__(self):
         assert self.contig is not None
-        if self.start is not None:
+        if self.start is None:
+            self.start = 1
+        else:
             self.start = int(self.start)
             assert self.start > 0
         if self.end is not None:

tests/test_simulated_data.py

@@ -1,6 +1,7 @@
 import sys
 
 import numpy.testing as nt
+import pysam
 import pytest
 import sgkit as sg
 
@@ -41,12 +42,13 @@ def assert_ts_ds_equal(ts, ds, ploidy=1):
     nt.assert_equal(ds.variant_position, ts.sites_position)
 
 
-def write_vcf(ts, vcf_path, contig_id="1"):
+def write_vcf(ts, vcf_path, contig_id="1", indexed=False):
     with open(vcf_path, "w") as f:
         ts.write_vcf(f, contig_id=contig_id)
-    # # This also compresses the input file
-    # pysam.tabix_index(str(vcf_path), preset="vcf")
-    # return vcf_path.with_suffix(vcf_path.suffix + ".gz")
+    if indexed:
+        # This also compresses the input file
+        pysam.tabix_index(str(vcf_path), preset="vcf")
+        vcf_path = vcf_path.with_suffix(vcf_path.suffix + ".gz")
     return vcf_path
 
 
@@ -75,6 +77,7 @@ def test_multi_contig(self, contig_ids, tmp_path):
             vcfs.append(vcf_path)
             tss[contig_id] = ts
 
+    def validate_tss_vcf_list(self, contig_ids, tss, vcfs, tmp_path):
         out = tmp_path / "example.vcf.zarr"
         vcf2zarr.convert(vcfs, out)
         ds = sg.load_dataset(out).set_index(
@@ -93,6 +96,34 @@ def test_multi_contig(self, contig_ids, tmp_path):
             dss = ds.sel(variants=(contig, slice(0, None)))
             assert_ts_ds_equal(tss[contig_id], dss)
 
+    @pytest.mark.parametrize("indexed", [True, False])
+    def test_indexed(self, indexed, tmp_path):
+        ts = run_simulation(num_samples=12, seed=34)
+        vcf_path = write_vcf(ts, tmp_path / "sim.vcf", indexed=indexed)
+        out = tmp_path / "example.vcf.zarr"
+        vcf2zarr.convert([vcf_path], out)
+        ds = sg.load_dataset(out)
+        assert_ts_ds_equal(ts, ds)
+
+    @pytest.mark.parametrize("num_contigs", [2, 3, 6])
+    def test_mixed_indexed(self, num_contigs, tmp_path):
+        contig_ids = [f"x{j}" for j in range(num_contigs)]
+
+        vcfs = []
+        tss = {}
+        for seed, contig_id in enumerate(contig_ids, 1):
+            ts = run_simulation(num_samples=3, seed=seed)
+            vcf_path = write_vcf(
+                ts,
+                tmp_path / f"{contig_id}.vcf",
+                contig_id=contig_id,
+                indexed=seed % 2 == 0,
+            )
+            vcfs.append(vcf_path)
+            tss[contig_id] = ts
+
+        self.validate_tss_vcf_list(contig_ids, tss, vcfs, tmp_path)
+
 
 # https://github.com/sgkit-dev/bio2zarr/issues/336
 @pytest.mark.skipif(sys.platform == "darwin", reason="msprime OSX pip packages broken")