OneScience-Group/CodonTransformer

1import tempfile2import unittest3 4import pandas as pd5from Bio.Data.CodonTable import TranslationError6 7from CodonTransformer.CodonData import (8    build_amino2codon_skeleton,9    get_amino_acid_sequence,10    is_correct_seq,11    preprocess_protein_sequence,12    read_fasta_file,13)14from CodonTransformer.CodonUtils import ProteinConfig15 16 17class TestCodonData(unittest.TestCase):18    def test_preprocess_protein_sequence(self):19        with ProteinConfig() as config:20            config.set("ambiguous_aminoacid_behavior", "raise_error")21            protein = "Z_"22            try:23                preprocess_protein_sequence(protein)24                self.fail("Expected ValueError")25            except ValueError:26                pass27            config.set("ambiguous_aminoacid_behavior", "standardize_deterministic")28            for _ in range(10):29                preprocessed_protein = preprocess_protein_sequence(protein)30                self.assertEqual(preprocessed_protein, "Q_")31            config.set("ambiguous_aminoacid_behavior", "standardize_random")32            random_results = set()33            # The probability of getting the same result 30 times in a row is34            # 1 in 1.073741824*10^9 if there are only two possible results.35            for _ in range(30):36                preprocessed_protein = preprocess_protein_sequence(protein)37                random_results.add(preprocessed_protein)38            self.assertGreater(len(random_results), 1)39 40    def test_read_fasta_file(self):41        fasta_content = ">sequence1\n" "ATGATGATGATGATG\n" ">sequence2\n" "TGATGATGATGA"42 43        with tempfile.NamedTemporaryFile(44            mode="w", delete=False, suffix=".fasta"45        ) as temp_file:46            temp_file.write(fasta_content)47            temp_file_name = temp_file.name48 49        try:50            sequences = read_fasta_file(temp_file_name, save_to_file=None)51            self.assertIsInstance(sequences, pd.DataFrame)52            self.assertEqual(len(sequences), 2)53            self.assertEqual(sequences.iloc[0]["dna"], "ATGATGATGATGATG")54            self.assertEqual(sequences.iloc[1]["dna"], "TGATGATGATGA")55        finally:56            import os57 58            os.unlink(temp_file_name)59 60    def test_build_amino2codon_skeleton(self):61        organism = "Homo sapiens"62        codon_skeleton = build_amino2codon_skeleton(organism)63 64        expected_amino_acids = "ARNDCQEGHILKMFPSTWYV_"65 66        for amino_acid in expected_amino_acids:67            self.assertIn(amino_acid, codon_skeleton)68            codons, frequencies = codon_skeleton[amino_acid]69            self.assertIsInstance(codons, list)70            self.assertIsInstance(frequencies, list)71            self.assertEqual(len(codons), len(frequencies))72            self.assertTrue(all(isinstance(codon, str) for codon in codons))73            self.assertTrue(all(freq == 0 for freq in frequencies))74 75        all_codons = set(76            codon for codons, _ in codon_skeleton.values() for codon in codons77        )78        self.assertEqual(len(all_codons), 64)  # There should be 64 unique codons79 80    def test_get_amino_acid_sequence(self):81        dna = "ATGGCCTGA"82        protein, is_correct = get_amino_acid_sequence(dna, return_correct_seq=True)83        self.assertEqual(protein, "MA_")84        self.assertTrue(is_correct)85 86    def test_is_correct_seq(self):87        dna = "ATGGCCTGA"88        protein = "MA_"89        self.assertTrue(is_correct_seq(dna, protein))90 91    def test_read_fasta_file_raises_exception_for_non_dna(self):92        non_dna_content = ">sequence1\nATGATGATGXYZATG\n>sequence2\nTGATGATGATGA"93 94        with tempfile.NamedTemporaryFile(95            mode="w", delete=False, suffix=".fasta"96        ) as temp_file:97            temp_file.write(non_dna_content)98            temp_file_name = temp_file.name99 100        try:101            with self.assertRaises(TranslationError) as context:102                read_fasta_file(temp_file_name)103            self.assertIn("Codon 'XYZ' is invalid", str(context.exception))104        finally:105            import os106 107            os.unlink(temp_file_name)108 109 110if __name__ == "__main__":111    unittest.main()112