finalize test

basf · Nov 20, 2024 · 8cfc27c · 8cfc27c
1 parent e279563
commit 8cfc27c
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/test_extras/test_chemprop/test_chemprop_pipeline.py b/test_extras/test_chemprop/test_chemprop_pipeline.py
@@ -308,7 +308,8 @@ def test_prediction(self) -> None:
         self.assertTrue(np.allclose(pred, pred_copy))
 
         # Test single prediction, this was causing an error before
-        _ = regression_model.predict([molecule_net_logd_df["smiles"].iloc[0]])
+        single_mol_pred = regression_model.predict([molecule_net_logd_df["smiles"].iloc[0]])
+        self.assertEqual(single_mol_pred.shape, (1,))
 
 
 class TestClassificationPipeline(unittest.TestCase):
@@ -345,7 +346,10 @@ def test_prediction(self) -> None:
         self.assertTrue(np.allclose(proba[~nan_indices], proba_copy[~nan_indices]))
 
         # Test single prediction, this was causing an error before
-        _ = classification_model.predict([molecule_net_bbbp_df["smiles"].iloc[0]])
+        single_mol_pred = classification_model.predict([molecule_net_bbbp_df["smiles"].iloc[0]])
+        self.assertEqual(single_mol_pred.shape, (1,))
+        single_mol_proba = classification_model.predict_proba([molecule_net_bbbp_df["smiles"].iloc[0]])
+        self.assertEqual(single_mol_proba.shape, (1, 2))
 
 
 class TestMulticlassClassificationPipeline(unittest.TestCase):