From c8b39636b7a079e018e971d189fc3bacebceb659 Mon Sep 17 00:00:00 2001
From: szhan <shing.zhan@gmail.com>
Date: Mon, 1 Jul 2024 12:26:27 +0100
Subject: [PATCH] Refactor tests for haploid case with multiallelic sites

---
 tests/lsbase.py                     | 34 ++------------------------
 tests/test_api_fb_haploid_multi.py  | 37 ++++-------------------------
 tests/test_api_vit_haploid_multi.py | 31 ++++--------------------
 3 files changed, 11 insertions(+), 91 deletions(-)

diff --git a/tests/lsbase.py b/tests/lsbase.py
index 402f08f..e0089b6 100644
--- a/tests/lsbase.py
+++ b/tests/lsbase.py
@@ -287,10 +287,10 @@ def get_ts_multiallelic_n10_no_recomb(self, seed=42):
         assert ts.num_sites > 3
         return ts
 
-    def get_ts_multiallelic_n6(self, seed=42):
+    def get_ts_multiallelic(self, num_samples, seed=42):
         ts = msprime.sim_mutations(
             msprime.sim_ancestry(
-                samples=6,
+                samples=num_samples,
                 recombination_rate=1e-4,
                 sequence_length=40,
                 population_size=1e4,
@@ -299,40 +299,10 @@ def get_ts_multiallelic_n6(self, seed=42):
             rate=1e-3,
             random_seed=seed,
         )
-        assert ts.num_sites > 5
-        return ts
-
-    def get_ts_multiallelic_n8(self, seed=42):
-        ts = msprime.sim_mutations(
-            msprime.sim_ancestry(
-                samples=8,
-                recombination_rate=1e-4,
-                sequence_length=20,
-                population_size=1e4,
-                random_seed=seed,
-            ),
-            rate=1e-4,
-            random_seed=seed,
-        )
         assert ts.num_trees > 15
         assert ts.num_sites > 5
         return ts
 
-    def get_ts_multiallelic_n16(self, seed=42):
-        ts = msprime.sim_mutations(
-            msprime.sim_ancestry(
-                samples=16,
-                recombination_rate=1e-2,
-                sequence_length=20,
-                population_size=1e4,
-                random_seed=seed,
-            ),
-            rate=1e-4,
-            random_seed=seed,
-        )
-        assert ts.num_sites > 5
-        return ts
-
 
 class ForwardBackwardAlgorithmBase(LSBase):
     """Base for testing forwards-backwards algorithms."""
diff --git a/tests/test_api_fb_haploid_multi.py b/tests/test_api_fb_haploid_multi.py
index 86ad692..34858b0 100644
--- a/tests/test_api_fb_haploid_multi.py
+++ b/tests/test_api_fb_haploid_multi.py
@@ -61,38 +61,11 @@ def test_ts_multiallelic_n10_no_recomb(
         self, scale_mutation_rate, include_ancestors
     ):
         ts = self.get_ts_multiallelic_n10_no_recomb()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
-
-    @pytest.mark.parametrize("scale_mutation_rate", [True, False])
-    @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n6(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n6()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
+        self.verify(ts, scale_mutation_rate, include_ancestors)
 
+    @pytest.mark.parametrize("num_samples", [6, 8, 16])
     @pytest.mark.parametrize("scale_mutation_rate", [True, False])
     @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n8(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n8()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
-
-    @pytest.mark.parametrize("scale_mutation_rate", [True, False])
-    @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n16(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n16()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
+    def test_ts_multiallelic_n16(self, num_samples, scale_mutation_rate, include_ancestors):
+        ts = self.get_ts_multiallelic(num_samples)
+        self.verify(ts, scale_mutation_rate, include_ancestors)
diff --git a/tests/test_api_vit_haploid_multi.py b/tests/test_api_vit_haploid_multi.py
index e93e853..4a55f14 100644
--- a/tests/test_api_vit_haploid_multi.py
+++ b/tests/test_api_vit_haploid_multi.py
@@ -50,32 +50,9 @@ def test_ts_multiallelic_n10_no_recomb(
             include_ancestors=include_ancestors,
         )
 
+    @pytest.mark.parametrize("num_samples", [6, 8, 16])
     @pytest.mark.parametrize("scale_mutation_rate", [True, False])
     @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n6(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n6()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
-
-    @pytest.mark.parametrize("scale_mutation_rate", [True, False])
-    @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n8(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n8()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
-
-    @pytest.mark.parametrize("scale_mutation_rate", [True, False])
-    @pytest.mark.parametrize("include_ancestors", [True, False])
-    def test_ts_multiallelic_n16(self, scale_mutation_rate, include_ancestors):
-        ts = self.get_ts_multiallelic_n16()
-        self.verify(
-            ts,
-            scale_mutation_rate=scale_mutation_rate,
-            include_ancestors=include_ancestors,
-        )
+    def test_ts_multiallelic_n16(self, num_samples, scale_mutation_rate, include_ancestors):
+        ts = self.get_ts_multiallelic(num_samples)
+        self.verify(ts, scale_mutation_rate, include_ancestors)