fix tests

Author: aditya0by0

chebai_proteins/preprocessing/datasets/deepGO/protein_pretraining.py

@@ -38,7 +38,7 @@ def __init__(self, **kwargs):
         Args:
             **kwargs: Additional arguments for the superclass initialization.
         """
-        self._go_uniprot_extractor = GOUniProtOver250()
+        self._go_uniprot_extractor = GOUniProtOver250(go_branch="all")
         assert self._go_uniprot_extractor.go_branch == GOUniProtOver250._ALL_GO_BRANCHES
 
         self.max_sequence_length: int = int(kwargs.get("max_sequence_length", 1002))
@@ -143,7 +143,6 @@ def _parse_protein_data_for_pretraining(self) -> pd.DataFrame:
             has_valid_associated_go_label = False
             for cross_ref in record.cross_references:
                 if cross_ref[0] == self._go_uniprot_extractor._GO_DATA_INIT:
-
                     if len(cross_ref) <= 3:
                         # No evidence code
                         continue

tests/unit/dataset_classes/testGOUniProDataExtractor.py

@@ -37,7 +37,7 @@ def setUpClass(
 
         _GOUniProtDataExtractor.READER = ProteinDataReader
 
-        cls.extractor = _GOUniProtDataExtractor()
+        cls.extractor = _GOUniProtDataExtractor(go_branch="all")
 
     def test_term_callback(self) -> None:
         """

tests/unit/dataset_classes/testGoUniProtOverX.py

@@ -17,7 +17,7 @@ def setUpClass(cls, mock_makedirs) -> None:
         """
         Set up the class for tests by initializing the extractor, graph, and input DataFrame.
         """
-        cls.extractor = _GOUniProtOverX()
+        cls.extractor = _GOUniProtOverX(go_branch="all")
         cls.test_graph: nx.DiGraph = GOUniProtMockData.get_transitively_closed_graph()
         cls.input_df: pd.DataFrame = GOUniProtMockData.get_data_in_dataframe().iloc[
             :, :4