Patch up tests for changed base_ts

Author: jeromekelleher

sc2ts/core.py

@@ -79,12 +79,16 @@ def get_problematic_sites():
 __cached_reference = None
 
 
-def get_reference_sequence():
+def get_reference_sequence(as_array=False):
     global __cached_reference
     if __cached_reference is None:
         reader = pyfaidx.Fasta(str(data_path / "reference.fasta"))
-        __cached_reference = "X" + str(reader[REFERENCE_GENBANK])
-    return __cached_reference
+        __cached_reference = reader[REFERENCE_GENBANK]
+    if as_array:
+        h = np.array(__cached_reference).astype(str)
+        return np.append(["X"], h)
+    else:
+        return "X" + str(__cached_reference)
 
 
 __cached_genes = None

tests/test_inference.py

@@ -148,7 +148,7 @@ def test_one_sample_one_mutation(self, tmp_path):
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
         assert ts2.num_mutations == 1
         var = next(ts2.variants())
-        assert var.alleles[var.genotypes[0]] == "X"
+        assert var.alleles[var.genotypes[1]] == "X"
 
     def test_one_sample_one_mutation_filtered(self, tmp_path):
         ts = sc2ts.initial_ts()
@@ -187,7 +187,7 @@ def test_two_samples_one_mutation_one_filtered(self, tmp_path):
         assert ts2.site(0).ancestral_state == ts.site(0).ancestral_state
         assert ts2.num_mutations == 1
         var = next(ts2.variants())
-        assert var.alleles[var.genotypes[0]] == "X"
+        assert var.alleles[var.genotypes[1]] == "X"
 
 
 class TestMatchTsinfer:
@@ -203,7 +203,7 @@ def test_match_reference(self, mirror):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         h = ma.alignment[ts.sites_position.astype(int)]
         samples[0].alignment = h
@@ -220,7 +220,7 @@ def test_match_reference_one_mutation(self, mirror, site_id):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         h = ma.alignment[ts.sites_position.astype(int)]
         # Mutate to gap
@@ -246,7 +246,7 @@ def test_match_reference_all_same(self, mirror, allele):
         tables.sites.truncate(20)
         ts = tables.tree_sequence()
         samples = util.get_samples(ts, [[(0, ts.sequence_length, 1)]])
-        alignment = sc2ts.core.get_reference_sequence()
+        alignment = sc2ts.core.get_reference_sequence(as_array=True)
         ma = sc2ts.alignments.encode_and_mask(alignment)
         ref = ma.alignment[ts.sites_position.astype(int)]
         h = np.zeros_like(ref) + allele
@@ -267,7 +267,7 @@ def match_path_ts(self, samples, ts):
         # FIXME this API is terrible
         ts2 = sc2ts.match_path_ts(samples, ts, samples[0].path, [])
         assert ts2.num_samples == len(samples)
-        for u, sample in zip(ts.samples(), samples):
+        for u, sample in zip(ts.samples()[1:], samples):
             node = ts.node(u)
             assert node.time == 0
             assert node.metadata == sample.metadata

tests/test_utils.py

@@ -9,7 +9,7 @@
 class TestPadSites:
     def check_site_padding(self, ts):
         ts = sc2ts.utils.pad_sites(ts)
-        ref = sc2ts.core.get_reference_sequence()
+        ref = sc2ts.core.get_reference_sequence(as_array=True)
         assert ts.num_sites == len(ref) - 1
         ancestral_state = ts.tables.sites.ancestral_state.view("S1").astype(str)
         assert np.all(ancestral_state == ref[1:])
@@ -67,7 +67,7 @@ def test_no_recombinants(self, ts, tmp_path):
 
     def test_one_sample_recombinant(self, tmp_path):
         ts = self.make_recombinant_tree(tmp_path / "match.db")
-        assert ts.num_samples == 3
+        assert ts.num_samples == 4
         re_nodes = [
             node.id for node in ts.nodes() if node.flags & sc2ts.NODE_IS_RECOMBINANT
         ]
@@ -93,6 +93,6 @@ def test_one_sample_recombinant(self, tmp_path):
     def test_two_sample_recombinant(self, tmp_path):
         """Test that we don't detach anything if the recombinant node is not a singleton"""
         ts = self.make_recombinant_tree(num_samples=2, db_path=tmp_path / "match.db")
-        assert ts.num_samples == 4
+        assert ts.num_samples == 5
         ts2 = utils.detach_singleton_recombinants(ts)
         ts.tables.assert_equals(ts2.tables, ignore_provenance=True)