Updated tests with new container types

EBjerrum · Mar 7, 2024 · 84d1315 · 84d1315
1 parent a6ac478
commit 84d1315
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -22,7 +22,9 @@ def smiles_list():
 
 _CONTAINER_CREATORS = [
     lambda x: x,
+    lambda x: np.array(x),
     lambda x: np.array(x).reshape(-1, 1),
+    lambda x: pd.Series(x),
     lambda x: pd.DataFrame({"hello": x}),
 ]
 

diff --git a/tests/test_smilestomol.py b/tests/test_smilestomol.py
@@ -30,10 +30,14 @@ def test_smilestomol_unsanitzable(invalid_smiles_list, smilestomol_transformer):
     with pytest.raises(ValueError):
         smilestomol_transformer.transform(invalid_smiles_list)
 
-def test_descriptor_transformer_parallel(smiles_list, smilestomol_transformer):
+def test_descriptor_transformer_parallel(smiles_container, smilestomol_transformer):
     smilestomol_transformer.set_params(parallel=True)
-    mol_list = smilestomol_transformer.transform(smiles_list)
-    assert all([ a == b for a, b in zip(smiles_list, [Chem.MolToSmiles(mol) for mol in mol_list.flatten()])])
+    mol_list = smilestomol_transformer.transform(smiles_container)
+    if isinstance(smiles_container, pd.DataFrame):
+        expected_smiles = smiles_container.iloc[:, 0].tolist()
+    else:
+        expected_smiles = smiles_container
+    assert all([ a == b for a, b in zip(expected_smiles, [Chem.MolToSmiles(mol) for mol in mol_list.flatten()])])
 
 def test_pandas_output(smiles_container, smilestomol_transformer, pandas_output):
         mols = smilestomol_transformer.transform(smiles_container)