from genomic_consultant.store.query import GenomicStore
from genomic_consultant.utils.models import FilterConfig, Variant


def test_filter_by_gene_and_af():
    store = GenomicStore(
        variants=[
            Variant(chrom="1", pos=1, ref="A", alt="T", gene="GENE1", allele_frequency=0.02),
            Variant(chrom="1", pos=2, ref="G", alt="C", gene="GENE2", allele_frequency=0.0001),
        ]
    )
    res = store.get_variants_by_gene("ind", ["GENE2"], filters=FilterConfig(max_af=0.001))
    assert len(res) == 1
    assert res[0].gene == "GENE2"


def test_consequence_include_exclude():
    store = GenomicStore(
        variants=[
            Variant(chrom="1", pos=1, ref="A", alt="T", gene="GENE1", consequence="missense_variant"),
            Variant(chrom="1", pos=2, ref="G", alt="C", gene="GENE1", consequence="synonymous_variant"),
        ]
    )
    res = store.get_variants_by_gene(
        "ind", ["GENE1"], filters=FilterConfig(consequence_includes=["missense"], consequence_excludes=["synonymous"])
    )
    assert len(res) == 1
    assert res[0].consequence == "missense_variant"