Tests for get_good_interpae.py

KosinskiLab · Sep 25, 2024 · 43ea8ec · 43ea8ec
1 parent 2b7c334
commit 43ea8ec
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 22 deletions.
diff --git a/.github/workflows/github_actions.yml b/.github/workflows/github_actions.yml
@@ -65,6 +65,7 @@ jobs:
           pytest -s test/test_features_with_templates.py
           pytest -s test/test_post_prediction.py
           pytest -s test/test_pdb_analyser.py
+          pytest -s test/test_get_good_inter_pae.py
 
   build-fold-container:
     runs-on: ubuntu-latest

diff --git a/alphapulldown/analysis_pipeline/get_good_inter_pae.py b/alphapulldown/analysis_pipeline/get_good_inter_pae.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
-from calculate_mpdockq import *
-from pdb_analyser import PDBAnalyser
+from .calculate_mpdockq import get_best_plddt, read_pdb, read_plddt, score_complex, calculate_mpDockQ, read_pdb_pdockq, calc_pdockq
+from .pdb_analyser import PDBAnalyser
 from Bio.PDB import PDBParser
 from Bio.PDB.Polypeptide import PPBuilder
 import os
@@ -23,12 +23,12 @@
 
 def examine_inter_pae(pae_mtx, seq_lengths, cutoff):
     """A function that checks inter-pae values in multimer prediction jobs"""
-    old_lenth = 0
+    old_length = 0
     mtx = pae_mtx.copy()
     for length in seq_lengths:
-        new_length = old_lenth + length
-        mtx[old_lenth:new_length, old_lenth:new_length] = 50
-        old_lenth = new_length
+        new_length = old_length + length
+        mtx[old_length:new_length, old_length:new_length] = 50
+        old_length = new_length
     check = np.where(mtx < cutoff)[0].size != 0
 
     return check

diff --git a/test/test_get_good_inter_pae.py b/test/test_get_good_inter_pae.py
@@ -1,15 +1,15 @@
 from absl.testing import absltest
 from unittest.mock import patch, mock_open, MagicMock
-import os
 import json
 import pandas as pd
 from alphapulldown.analysis_pipeline.get_good_inter_pae import obtain_seq_lengths, main
 
+
 class TestGetGoodInterPae(absltest.TestCase):
 
-    @patch('get_good_inter_pae.os.path.exists')
-    @patch('get_good_inter_pae.PDBParser.get_structure')
-    @patch('get_good_inter_pae.PPBuilder.build_peptides')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.os.path.exists')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.PDBParser.get_structure')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.PPBuilder.build_peptides')
     def test_obtain_seq_lengths(self, mock_build_peptides, mock_get_structure, mock_exists):
         mock_exists.return_value = True
         mock_get_structure.return_value = MagicMock()
@@ -18,28 +18,39 @@ def test_obtain_seq_lengths(self, mock_build_peptides, mock_get_structure, mock_
         result = obtain_seq_lengths('/fake/dir')
         self.assertEqual(result, [4])
 
-    @patch('get_good_inter_pae.os.listdir')
-    @patch('get_good_inter_pae.os.path.isfile')
-    @patch('get_good_inter_pae.PDBAnalyser')
-    @patch('get_good_inter_pae.obtain_pae_and_iptm')
-    @patch('get_good_inter_pae.obtain_seq_lengths')
-    @patch('get_good_inter_pae.examine_inter_pae')
-    @patch('get_good_inter_pae.obtain_mpdockq')
-    @patch('get_good_inter_pae.pd.DataFrame.to_csv')
-    def test_main(self, mock_to_csv, mock_obtain_mpdockq, mock_examine_inter_pae, mock_obtain_seq_lengths, mock_obtain_pae_and_iptm, mock_PDBAnalyser, mock_isfile, mock_listdir):
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.os.listdir')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.os.path.isfile')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.PDBAnalyser')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.obtain_pae_and_iptm')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.obtain_seq_lengths')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.examine_inter_pae')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.obtain_mpdockq')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.pd.DataFrame.to_csv')
+    @patch('alphapulldown.analysis_pipeline.get_good_inter_pae.os.makedirs')
+    def test_main(self, mock_makedirs, mock_to_csv, mock_obtain_mpdockq, mock_examine_inter_pae,
+                  mock_obtain_seq_lengths, mock_obtain_pae_and_iptm, mock_PDBAnalyser, mock_isfile, mock_listdir):
         mock_listdir.return_value = ['job1']
         mock_isfile.return_value = True
         mock_obtain_pae_and_iptm.return_value = (MagicMock(), 0.5)
         mock_obtain_seq_lengths.return_value = [100, 100]
         mock_examine_inter_pae.return_value = True
         mock_obtain_mpdockq.return_value = (0.5, {'A': [0.9] * 100, 'B': [0.9] * 100})
-        mock_PDBAnalyser.return_value = MagicMock()
 
-        with patch('builtins.open', mock_open(read_data=json.dumps({'order': ['model1'], 'iptm+ptm': {'model1': 0.8}}))):
-            with patch('get_good_inter_pae.FLAGS', MagicMock(output_dir='/fake/dir', cutoff=5)):
+        # Mocking PDBAnalyser to return a DataFrame with expected columns
+        mock_PDBAnalyser.return_value = MagicMock(return_value=pd.DataFrame({
+            'pdb': ['pdb1'],
+            'pvalue': [0.01],
+            'Hydrophobhic': [0.5]
+        }))
+
+        with patch('builtins.open',
+                   mock_open(read_data=json.dumps({'order': ['model1'], 'iptm+ptm': {'model1': 0.8}}))):
+            with patch('alphapulldown.analysis_pipeline.get_good_inter_pae.FLAGS',
+                       MagicMock(output_dir='/fake/dir', cutoff=5)):
                 main([])
 
         mock_to_csv.assert_called_once()
 
+
 if __name__ == '__main__':
     absltest.main()