Add tests

Author: fellen31

genmod/annotate_models/models/compound_model.py

@@ -12,13 +12,11 @@
 
 from __future__ import print_function
 from genmod.vcf_tools.genotype import Genotype
-from intervaltree import IntervalTree
 import logging
 from typing import Union
 
-# TODO: What is family?
-# TODO: Remove intervals
-def check_compounds(variant_1: dict, variant_2: dict, family, intervals: IntervalTree, phased: bool) -> bool:
+
+def check_compounds(variant_1: dict, variant_2: dict, family, phased: bool) -> bool:
     """
     Check if two variants of a pair follow the compound heterozygous model.
 
@@ -56,8 +54,12 @@ def check_compounds(variant_1: dict, variant_2: dict, family, intervals: Interva
             for parent_id in (individual.mother, individual.father):
                 if parent_id != "0":
                     parent = family.individuals[parent_id]
-                    parent_genotypes = [get_genotype(variant, parent_id) for variant in (variant_1, variant_2)]
-                    if parent.healthy and all(genotype.has_variant for genotype in parent_genotypes):
+                    parent_genotypes = [
+                        get_genotype(variant, parent_id) for variant in (variant_1, variant_2)
+                    ]
+                    if parent.healthy and all(
+                        genotype.has_variant for genotype in parent_genotypes
+                    ):
                         return False
 
         genotype_1 = get_genotype(variant_1, individual_id)
@@ -69,13 +71,17 @@ def check_compounds(variant_1: dict, variant_2: dict, family, intervals: Interva
             return False
 
         # If the individual is affected and phased we can say that the variants need to be on different alleles, otherwise it can not be a compound pair.
-        if individual.affected and phased and variants_on_same_allele(individual.individual_id, variant_1, variant_2):
+        if (
+            individual.affected
+            and phased
+            and variants_on_same_allele(individual.individual_id, variant_1, variant_2)
+        ):
             return False
 
     # If the individual is affected and not phased we can not say anything about the phase, so we say it is a compound pair, since it could be a compound pair.
     return True
 
-# TODO: Write a test for this
+
 def get_genotype(variant: dict, individual_id: str) -> Genotype:
     """
     Return the Genotype object for a variants for a given individual.
@@ -89,7 +95,7 @@ def get_genotype(variant: dict, individual_id: str) -> Genotype:
     """
     return variant["genotypes"][individual_id]
 
-# TODO: Write a test for this
+
 def get_phase_set(variant_dict: dict, sample_id: str) -> Union[str, None]:
     """
     Extracts the PS (phase set) field for a given sample from a VCF-like variant dictionary.
@@ -106,7 +112,7 @@ def get_phase_set(variant_dict: dict, sample_id: str) -> Union[str, None]:
     field_map = dict(zip(format_fields, sample_values))
     return field_map.get("PS")
 
-# TODO: Write a test for this
+
 def variants_on_same_allele(individual_id: str, variant_1: dict, variant_2: dict) -> bool:
     """
     Determine if two phased variants are on the same haplotype for a given individual.
@@ -130,13 +136,13 @@ def variants_on_same_allele(individual_id: str, variant_1: dict, variant_2: dict
 
     same_phase = get_phase_set(variant_1, individual_id) == get_phase_set(variant_2, individual_id)
     overlapping_allele = (
-        genotype_1.allele_1 == genotype_2.allele_1 or
-        genotype_1.allele_2 == genotype_2.allele_2
+        genotype_1.allele_1 == genotype_2.allele_1 or genotype_1.allele_2 == genotype_2.allele_2
     )
 
     # Variants are on the same allele if they overlap in the same haplotype and are in the same phase set
     return same_phase and overlapping_allele
 
+
 def main():
     pass
 

tests/genetic_models/test_compound_model.py

@@ -0,0 +1,139 @@
+from ped_parser import FamilyParser
+
+from genmod.annotate_models.models.compound_model import (
+    check_compounds,
+    get_genotype,
+    get_phase_set,
+    variants_on_same_allele,
+)
+from genmod.vcf_tools import Genotype
+
+
+def get_family(family_lines):
+    return FamilyParser(family_lines)
+
+
+def make_variant(*, genotypes, format_str="GT:PS", sample_fields=None):
+    variant = {"genotypes": dict(genotypes), "FORMAT": format_str}
+    if sample_fields:
+        variant.update(sample_fields)
+    return variant
+
+
+def test_get_genotype_returns_expected_object():
+    genotype = Genotype(**{"GT": "0/1"})
+    variant = {"genotypes": {"sample": genotype}}
+    assert get_genotype(variant, "sample") is genotype
+
+
+def test_get_phase_set_returns_ps_when_present():
+    variant = {"FORMAT": "GT:AD:PS", "sample": "0/1:10,5:12345"}
+    assert get_phase_set(variant, "sample") == "12345"
+
+
+def test_get_phase_set_returns_none_when_ps_missing_in_format():
+    variant = {"FORMAT": "GT:AD", "sample": "0/1:10,5"}
+    assert get_phase_set(variant, "sample") is None
+
+
+def test_get_phase_set_returns_none_when_sample_missing_or_unset():
+    variant_missing_sample = {"FORMAT": "GT:PS"}
+    assert get_phase_set(variant_missing_sample, "sample") is None
+
+    variant_empty_sample = {"FORMAT": "GT:PS", "sample": ""}
+    assert get_phase_set(variant_empty_sample, "sample") is None
+
+    variant_short_sample = {"FORMAT": "GT:PS", "sample": "0|1"}
+    assert get_phase_set(variant_short_sample, "sample") is None
+
+
+def test_variants_on_same_allele_true_same_ps_and_overlap():
+    variant_1 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "0|1"})},
+        sample_fields={"sample": "0|1:PS1"},
+    )
+    variant_2 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "0|1"})},
+        sample_fields={"sample": "0|1:PS1"},
+    )
+    assert variants_on_same_allele("sample", variant_1, variant_2) is True
+
+
+def test_variants_on_same_allele_false_different_phase_set():
+    variant_1 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "0|1"})},
+        sample_fields={"sample": "0|1:1"},
+    )
+    variant_2 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "0|1"})},
+        sample_fields={"sample": "0|1:2"},
+    )
+    assert variants_on_same_allele("sample", variant_1, variant_2) is False
+
+
+def test_variants_on_same_allele_false_same_phase_set_no_overlap():
+    variant_1 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "0|1"})},
+        sample_fields={"sample": "0|1:1"},
+    )
+    variant_2 = make_variant(
+        genotypes={"sample": Genotype(**{"GT": "1|0"})},
+        sample_fields={"sample": "1|0:1"},
+    )
+    assert variants_on_same_allele("sample", variant_1, variant_2) is False
+
+
+def test_check_compounds_false_when_affected_phased_same_allele():
+    family_lines = [
+        "#FamilyID\tSampleID\tFather\tMother\tSex\tPhenotype\n",
+        "1\tchild\tfather\tmother\t1\t2\n",
+        "1\tfather\t0\t0\t1\t1\n",
+        "1\tmother\t0\t0\t2\t1\n",
+    ]
+    family = get_family(family_lines)
+
+    variant_1 = make_variant(
+        genotypes={
+            "child": Genotype(**{"GT": "0|1"}),
+            "father": Genotype(**{"GT": "0|0"}),
+            "mother": Genotype(**{"GT": "0|0"}),
+        },
+        sample_fields={"child": "0|1:1"},
+    )
+    variant_2 = make_variant(
+        genotypes={
+            "child": Genotype(**{"GT": "0|1"}),
+            "father": Genotype(**{"GT": "0|0"}),
+            "mother": Genotype(**{"GT": "0|0"}),
+        },
+        sample_fields={"child": "0|1:1"},
+    )
+
+    assert check_compounds(variant_1, variant_2, family=family, phased=True) is False
+
+
+def test_check_compounds_false_when_healthy_parent_has_both_variants():
+    family_lines = [
+        "#FamilyID\tSampleID\tFather\tMother\tSex\tPhenotype\n",
+        "1\tchild\tfather\tmother\t1\t2\n",
+        "1\tfather\t0\t0\t1\t1\n",
+        "1\tmother\t0\t0\t2\t1\n",
+    ]
+    family = get_family(family_lines)
+
+    variant_1 = make_variant(
+        genotypes={
+            "child": Genotype(**{"GT": "0/1"}),
+            "father": Genotype(**{"GT": "0/1"}),
+            "mother": Genotype(**{"GT": "0/0"}),
+        }
+    )
+    variant_2 = make_variant(
+        genotypes={
+            "child": Genotype(**{"GT": "0/1"}),
+            "father": Genotype(**{"GT": "0/1"}),
+            "mother": Genotype(**{"GT": "0/0"}),
+        }
+    )
+
+    assert check_compounds(variant_1, variant_2, family=family, phased=False) is False