test_libsnp.py

#!/usr/bin/env python3
# Copyright (C) 2023 Blaise Li
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
Tests for libsnp.
"""

from pathlib import Path
import sys
import unittest

test_dir = Path(__file__).parent
lib_dir = test_dir.parent.joinpath("src")
sys.path.insert(0, str(lib_dir))

import pandas as pd
from libsnp import load_counts_from_joint_vcf, concat_counts_from_dicts


class TestLoader(unittest.TestCase):
    """Tests for the counts loader."""
    def setUp(self):
        self.chroms = [str(i) for i in range(1, 37)] + ["maxi"]
        self.test_data = test_dir.joinpath("data", "jointcall_sample.vcf.gz")
        self.counts_dict = load_counts_from_joint_vcf(
            self.test_data,
            chroms=self.chroms)
        self.samples = list(self.counts_dict.keys())

    def test_load_data(self):
        """Some iintegrity checks on loaded table."""
        self.assertEqual(len(self.samples), 23)

    def test_concat_counts_tables(self):
        """Check that count tables concatenate."""
        all_counts = pd.concat(
            self.counts_dict, axis=1)
        (_, ncols) = all_counts.shape
        self.assertEqual(ncols, 2 * len(self.samples))

class TestConcatenator(unittest.TestCase):
    """Tests for the counts loader."""
    def setUp(self):
        self.chroms = [str(i) for i in range(1, 37)] + ["maxi"]
        self.test_data = test_dir.joinpath("data", "jointcall_sample.vcf.gz")
        self.counts_dict = load_counts_from_joint_vcf(
            self.test_data,
            chroms=self.chroms)
        self.samples = list(self.counts_dict.keys())

    def test_concat_data(self):
        """Some iintegrity checks on loaded table."""
        counts = concat_counts_from_dicts(self.counts_dict)
        self.assertEqual(counts.shape[1], len(self.samples) * 2)

if __name__ == "__main__":
    unittest.main()