From e936616d6c87d03106dc99d8dbf823feedaead4c Mon Sep 17 00:00:00 2001 From: Brian Yee Date: Sun, 4 Oct 2026 09:31:07 -0700 Subject: [PATCH] test: cover annotation lookup helpers --- .github/workflows/unit-tests.yml | 18 ++++ annotator/annotation_functions.py | 45 ++++----- unit_tests/test_annotation_functions.py | 121 ++++++++++++++++++++++++ 3 files changed, 157 insertions(+), 27 deletions(-) create mode 100644 .github/workflows/unit-tests.yml create mode 100644 unit_tests/test_annotation_functions.py diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml new file mode 100644 index 0000000..48ad792 --- /dev/null +++ b/.github/workflows/unit-tests.yml @@ -0,0 +1,18 @@ +name: Unit tests + +on: + pull_request: + branches: [master] + +jobs: + unit-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + - name: Install test dependencies + run: python -m pip install pytest gffutils pybedtools pandas numpy future tqdm + - name: Run unit tests + run: python -m pytest -q unit_tests diff --git a/annotator/annotation_functions.py b/annotator/annotation_functions.py index c3d099a..c2068a2 100644 --- a/annotator/annotation_functions.py +++ b/annotator/annotation_functions.py @@ -132,14 +132,12 @@ def get_keys(gtf_format): def get_gene_to_transcript_dict(db, gene_id_key, transcript_id_key): """ - Returns a gene: transcript dictionary. - - gene:[transcript1, transcript2] - - :param db: gffutils.FeatureDB - :param gene_id_key: string - :param transcript_id_key: string - :return: + Map each gene ID to the transcript IDs on its transcript features. + + :param db: gffutils.FeatureDB containing transcript features + :param gene_id_key: attribute name containing gene IDs + :param transcript_id_key: attribute name containing transcript IDs + :return: defaultdict(list) mapping gene IDs to lists of transcript IDs """ genes_dict = defaultdict(list) @@ -234,18 +232,12 @@ def get_all_cds_dict(db, cds_key): def get_all_exons_dict(db, exon_key, transcript_id_key): """ - Returns dictionary of exons as transcript_id:{ - [ - {'start':START, 'end':END}, - {'start':START, 'end':END}, - ... - ] - }. - - :param db: gffutils.FeatureDB - :param transcript_id_key: string - :param exon_key: string - :return: + Collect exon coordinates, chromosome, and strand by transcript ID. + + :param db: gffutils.FeatureDB containing exon features + :param exon_key: feature type to treat as an exon + :param transcript_id_key: attribute name containing transcript IDs + :return: defaultdict(list) mapping transcript IDs to exon dictionaries """ exons_dict = defaultdict(list) @@ -264,13 +256,12 @@ def get_all_exons_dict(db, exon_key, transcript_id_key): def get_all_transcripts_dict(db, transcript_key, transcript_id_key): """ - Returns dictionary of transcript_id:{'start':START, 'end':END}. - - :param db: gffutils.FeatureDB - :param transcript_key: string - :param transcript_id_key: string - :return: chr19_transcripts_dict: defaultdict(dict) - hash of transcripts and their start/end coordinates + Collect the start and end coordinates of each transcript ID. + + :param db: gffutils.FeatureDB containing transcript features + :param transcript_key: feature type to treat as a transcript + :param transcript_id_key: attribute name containing transcript IDs + :return: defaultdict(dict) mapping transcript IDs to start/end dictionaries """ transcripts_dict = defaultdict(dict) for transcript_feature in db.features_of_type(transcript_key): diff --git a/unit_tests/test_annotation_functions.py b/unit_tests/test_annotation_functions.py new file mode 100644 index 0000000..df3d993 --- /dev/null +++ b/unit_tests/test_annotation_functions.py @@ -0,0 +1,121 @@ +"""Unit tests for the public annotation lookup helpers.""" + +from types import SimpleNamespace + +from annotator.annotation_functions import ( + gene_id_to_name, + get_all_exons_dict, + get_all_transcripts_dict, + get_gene_to_transcript_dict, +) + + +class FeatureDBStub: + def __init__(self, features): + self.features = features + + def features_of_type(self, feature_type): + return iter(self.features.get(feature_type, ())) + + +def feature(attributes, start=1, end=10, seqid="chr1", strand="+"): + return SimpleNamespace( + attributes=attributes, + start=start, + end=end, + seqid=seqid, + strand=strand, + ) + + +def test_gene_to_transcript_dict_groups_transcripts_by_gene(): + db = FeatureDBStub( + { + "transcript": [ + feature({"gene_id": ["gene-a"], "transcript_id": ["tx-1"]}), + feature({"gene_id": ["gene-a"], "transcript_id": ["tx-2"]}), + feature({"gene_id": ["gene-b"], "transcript_id": ["tx-3"]}), + ], + "gene": [feature({"gene_id": ["unrelated"]})], + } + ) + + assert get_gene_to_transcript_dict(db, "gene_id", "transcript_id") == { + "gene-a": ["tx-1", "tx-2"], + "gene-b": ["tx-3"], + } + + +def test_gene_to_transcript_dict_uses_supplied_attribute_keys(): + db = FeatureDBStub( + { + "transcript": [ + feature({"gene": ["gene-a", "gene-b"], "isoform": ["tx-1", "tx-2"]}) + ] + } + ) + + assert get_gene_to_transcript_dict(db, "gene", "isoform") == { + "gene-a": ["tx-1", "tx-2"], + "gene-b": ["tx-1", "tx-2"], + } + + +def test_gene_to_transcript_dict_returns_empty_mapping_without_transcripts(): + db = FeatureDBStub({"gene": [feature({"gene_id": ["gene-a"]})]}) + + assert get_gene_to_transcript_dict(db, "gene_id", "transcript_id") == {} + + +def test_all_transcripts_dict_preserves_coordinates_for_each_transcript_id(): + db = FeatureDBStub( + { + "mRNA": [ + feature({"ID": ["tx-1", "tx-2"]}, start=10, end=30), + feature({"ID": ["tx-3"]}, start=40, end=70), + ], + "transcript": [feature({"ID": ["unrelated"]})], + } + ) + + assert get_all_transcripts_dict(db, "mRNA", "ID") == { + "tx-1": {"start": 10, "end": 30}, + "tx-2": {"start": 10, "end": 30}, + "tx-3": {"start": 40, "end": 70}, + } + + +def test_all_exons_dict_keeps_each_exon_and_its_strand(): + db = FeatureDBStub( + { + "exon": [ + feature({"transcript_id": ["tx-1", "tx-2"]}, 10, 20, "chr2", "-"), + feature({"transcript_id": ["tx-1"]}, 30, 40, "chr2", "-"), + ] + } + ) + + assert get_all_exons_dict(db, "exon", "transcript_id") == { + "tx-1": [ + {"chrom": "chr2", "start": 10, "end": 20, "strand": "-"}, + {"chrom": "chr2", "start": 30, "end": 40, "strand": "-"}, + ], + "tx-2": [{"chrom": "chr2", "start": 10, "end": 20, "strand": "-"}], + } + + +def test_gene_id_to_name_accepts_list_and_scalar_gene_ids(): + db = FeatureDBStub( + { + "gene": [ + feature({"gene_id": ["gene-a"], "symbol": ["Alpha"]}), + feature({"gene_id": "gene-b", "symbol": ["Beta"]}), + ], + "transcript": [feature({"gene_id": ["unrelated"], "symbol": ["Other"]})], + } + ) + + assert gene_id_to_name(db, "symbol") == { + "gene-a": "Alpha", + "gene-b": "Beta", + }