OneScience-Group/CodonTransformer

1import unittest2 3import ipywidgets4 5from CodonTransformer.CodonJupyter import (6    DNASequencePrediction,7    UserContainer,8    create_dropdown_options,9    create_organism_dropdown,10    display_organism_dropdown,11    display_protein_input,12    format_model_output,13)14from CodonTransformer.CodonUtils import ORGANISM2ID15 16 17class TestCodonJupyter(unittest.TestCase):18    def test_UserContainer(self):19        user_container = UserContainer()20        self.assertEqual(user_container.organism, -1)21        self.assertEqual(user_container.protein, "")22 23    def test_create_organism_dropdown(self):24        container = UserContainer()25        dropdown = create_organism_dropdown(container)26 27        self.assertIsInstance(dropdown, ipywidgets.Dropdown)28        self.assertGreater(len(dropdown.options), 0)29        self.assertEqual(dropdown.description, "")30        self.assertEqual(dropdown.layout.width, "40%")31        self.assertEqual(dropdown.layout.margin, "0 0 10px 0")32        self.assertEqual(dropdown.style.description_width, "initial")33 34        # Test the dropdown options35        options = dropdown.options36        self.assertIn("", options)37        self.assertIn("Selected Organisms", options)38        self.assertIn("All Organisms", options)39 40    def test_create_dropdown_options(self):41        options = create_dropdown_options(ORGANISM2ID)42        self.assertIsInstance(options, list)43        self.assertGreater(len(options), 0)44 45    def test_display_organism_dropdown(self):46        container = UserContainer()47        with unittest.mock.patch(48            "CodonTransformer.CodonJupyter.display"49        ) as mock_display:50            display_organism_dropdown(container)51 52        # Check that display was called twice (for container_widget and HTML)53        self.assertEqual(mock_display.call_count, 2)54 55        # Check that the first call to display was with a VBox widget56        self.assertIsInstance(mock_display.call_args_list[0][0][0], ipywidgets.VBox)57 58        # Check that the VBox contains a Dropdown59        dropdown = mock_display.call_args_list[0][0][0].children[1]60        self.assertIsInstance(dropdown, ipywidgets.Dropdown)61        self.assertGreater(len(dropdown.options), 0)62 63    def test_display_protein_input(self):64        container = UserContainer()65        with unittest.mock.patch(66            "CodonTransformer.CodonJupyter.display"67        ) as mock_display:68            display_protein_input(container)69 70        # Check that display was called twice (for container_widget and HTML)71        self.assertEqual(mock_display.call_count, 2)72 73        # Check that the first call to display was with a VBox widget74        self.assertIsInstance(mock_display.call_args_list[0][0][0], ipywidgets.VBox)75 76        # Check that the VBox contains a Textarea77        textarea = mock_display.call_args_list[0][0][0].children[1]78        self.assertIsInstance(textarea, ipywidgets.Textarea)79 80        # Verify the properties of the Textarea81        self.assertEqual(textarea.value, "")82        self.assertEqual(textarea.placeholder, "Enter here...")83        self.assertEqual(textarea.description, "")84        self.assertEqual(textarea.layout.width, "100%")85        self.assertEqual(textarea.layout.height, "100px")86        self.assertEqual(textarea.layout.margin, "0 0 10px 0")87        self.assertEqual(textarea.style.description_width, "initial")88 89    def test_format_model_output(self):90        output = DNASequencePrediction(91            organism="Escherichia coli",92            protein="MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",93            processed_input="MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",94            predicted_dna="ATGAAAACTGTTCGTCAGGAACGTCTGAAATCTATTGTTCGTATTCTGGAACGTTCTAAAGAACCGGTTTCTGGTGCTCAACTGGCTGAAGAACTGTCTGTTTCTCGTCAGGTTATTGTTCAGGACATTGCTTACCTGCGTTCTCTGGGTTATAA",95        )96        formatted_output = format_model_output(output)97        self.assertIsInstance(formatted_output, str)98        self.assertIn("Organism", formatted_output)99        self.assertIn("Escherichia coli", formatted_output)100        self.assertIn("Input Protein", formatted_output)101        self.assertIn(102            "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",103            formatted_output,104        )105        self.assertIn("Processed Input", formatted_output)106        self.assertIn(107            "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG",108            formatted_output,109        )110        self.assertIn("Predicted DNA", formatted_output)111        self.assertIn(112            "ATGAAAACTGTTCGTCAGGAACGTCTGAAATCTATTGTTCGTATTCTGGAACGTTCTAAAGAACCGGTTTCTGGTGCTCAACTGGCTGAAGAACTGTCTGTTTCTCGTCAGGTTATTGTTCAGGACATTGCTTACCTGCGTTCTCTGGGTTATAA",113            formatted_output,114        )115 116 117if __name__ == "__main__":118    unittest.main()119