OneScience-Group/CodonTransformer

1import os2import pickle3import tempfile4import unittest5 6from CodonTransformer.CodonUtils import (7    ProteinConfig,8    find_pattern_in_fasta,9    get_organism2id_dict,10    get_taxonomy_id,11    load_pkl_from_url,12    load_python_object_from_disk,13    save_python_object_to_disk,14    sort_amino2codon_skeleton,15)16 17 18class TestCodonUtils(unittest.TestCase):19    def test_config_manager(self):20        with ProteinConfig() as config:21            config.set("ambiguous_aminoacid_behavior", "standardize_deterministic")22            self.assertEqual(23                config.get("ambiguous_aminoacid_behavior"), "standardize_deterministic"24            )25            config.set("ambiguous_aminoacid_map_override", {"X": ["A", "G"]})26            self.assertEqual(27                config.get("ambiguous_aminoacid_map_override"), {"X": ["A", "G"]}28            )29            config.update(30                {31                    "ambiguous_aminoacid_behavior": "raise_error",32                    "ambiguous_aminoacid_map_override": {"X": ["A", "G"]},33                }34            )35            self.assertEqual(config.get("ambiguous_aminoacid_behavior"), "raise_error")36            self.assertEqual(37                config.get("ambiguous_aminoacid_map_override"), {"X": ["A", "G"]}38            )39            try:40                config.set("invalid_key", "invalid_value")41                self.fail("Expected ValueError")42            except ValueError:43                pass44        with ProteinConfig() as config:45            self.assertEqual(46                config.get("ambiguous_aminoacid_behavior"), "standardize_random"47            )48            self.assertEqual(config.get("ambiguous_aminoacid_map_override"), {})49 50    def test_load_python_object_from_disk(self):51        test_obj = {"key1": "value1", "key2": 2}52        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as temp_file:53            temp_file_name = temp_file.name54            save_python_object_to_disk(test_obj, temp_file_name)55        loaded_obj = load_python_object_from_disk(temp_file_name)56        self.assertEqual(test_obj, loaded_obj)57        os.remove(temp_file_name)58 59    def test_save_python_object_to_disk(self):60        test_obj = [1, 2, 3, 4, 5]61        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as temp_file:62            temp_file_name = temp_file.name63            save_python_object_to_disk(test_obj, temp_file_name)64            self.assertTrue(os.path.exists(temp_file_name))65        os.remove(temp_file_name)66 67    def test_find_pattern_in_fasta(self):68        text = (69            ">seq1 [keyword=value1]\nATGCGTACGTAGCTAG\n"70            ">seq2 [keyword=value2]\nGGTACGATCGATCGAT"71        )72        self.assertEqual(find_pattern_in_fasta("keyword", text), "value1")73        self.assertEqual(find_pattern_in_fasta("nonexistent", text), "")74 75    def test_get_organism2id_dict(self):76        with tempfile.NamedTemporaryFile(77            mode="w", delete=True, suffix=".csv"78        ) as temp_file:79            temp_file.write("0,Escherichia coli\n1,Homo sapiens\n2,Mus musculus")80            temp_file.flush()81            organism2id = get_organism2id_dict(temp_file.name)82            self.assertEqual(83                organism2id,84                {"Escherichia coli": 0, "Homo sapiens": 1, "Mus musculus": 2},85            )86 87    def test_get_taxonomy_id(self):88        taxonomy_dict = {89            "Escherichia coli": 562,90            "Homo sapiens": 9606,91            "Mus musculus": 10090,92        }93        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=True) as temp_file:94            temp_file_name = temp_file.name95            save_python_object_to_disk(taxonomy_dict, temp_file_name)96            self.assertEqual(get_taxonomy_id(temp_file_name, "Escherichia coli"), 562)97            self.assertEqual(98                get_taxonomy_id(temp_file_name, return_dict=True), taxonomy_dict99            )100 101    def test_sort_amino2codon_skeleton(self):102        amino2codon = {103            "A": (["GCT", "GCC", "GCA", "GCG"], [0.0, 0.0, 0.0, 0.0]),104            "C": (["TGT", "TGC"], [0.0, 0.0]),105        }106        sorted_amino2codon = sort_amino2codon_skeleton(amino2codon)107        self.assertEqual(108            sorted_amino2codon,109            {110                "A": (["GCA", "GCC", "GCG", "GCT"], [0.0, 0.0, 0.0, 0.0]),111                "C": (["TGC", "TGT"], [0.0, 0.0]),112            },113        )114 115    def test_load_pkl_from_url(self):116        url = "https://example.com/test.pkl"117        expected_obj = {"key": "value"}118        with unittest.mock.patch("requests.get") as mock_get:119            mock_get.return_value.content = pickle.dumps(expected_obj)120            loaded_obj = load_pkl_from_url(url)121        self.assertEqual(loaded_obj, expected_obj)122 123 124if __name__ == "__main__":125    unittest.main()126