forked from enviPath/enviPy
92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
from unittest.mock import Mock, patch
|
|
|
|
from django.test import TestCase
|
|
|
|
from epdb.logic import SEdge, SNode, SPathway
|
|
from epdb.models import Pathway, Setting
|
|
from utilities.chem import PredictionResult, ProductSet
|
|
|
|
|
|
class SNodeTest(TestCase):
|
|
def test_snode_eq(self):
|
|
snode1 = SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)
|
|
snode2 = SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)
|
|
self.assertEqual(snode1, snode2)
|
|
|
|
|
|
class SEdgeTest(TestCase):
|
|
def test_sedge_eq(self):
|
|
sedge1 = SEdge(
|
|
[SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)],
|
|
[SNode("CN1C(=O)NC2=C(C1=O)N(C)C=N2", 1), SNode("C=O", 1)],
|
|
rule=None,
|
|
)
|
|
sedge2 = SEdge(
|
|
[SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)],
|
|
[SNode("CN1C(=O)NC2=C(C1=O)N(C)C=N2", 1), SNode("C=O", 1)],
|
|
rule=None,
|
|
)
|
|
self.assertEqual(sedge1, sedge2)
|
|
|
|
|
|
class SPathwayTest(TestCase):
|
|
def setUp(self):
|
|
"""Set up test data for SPathway tests."""
|
|
self.test_smiles = "CCN(CC)C(=O)C1=CC(=CC=C1)CO"
|
|
self.mock_setting = Mock(spec=Setting)
|
|
self.mock_pathway = Mock(spec=Pathway)
|
|
|
|
def test_predict_step_basic(self):
|
|
"""Test basic predict_step functionality."""
|
|
spw = SPathway(root_nodes=self.test_smiles, prediction_setting=self.mock_setting)
|
|
|
|
# e.g. bt0002
|
|
pr = PredictionResult(
|
|
[
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(=O)C(O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(O)C2=O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(=O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
],
|
|
0.17,
|
|
None,
|
|
)
|
|
|
|
mock_val = {"rule_triggered": True, "transformations": [pr]}
|
|
|
|
with patch.object(self.mock_setting, "expand", return_value=mock_val):
|
|
spw.predict_step(from_depth=0)
|
|
|
|
self.assertEqual(len(spw.smiles_to_node.keys()), 4)
|
|
self.assertEqual(len(spw.edges), 3)
|
|
|
|
def test_to_json(self):
|
|
"""Test basic predict_step functionality."""
|
|
spw = SPathway(root_nodes=self.test_smiles, prediction_setting=self.mock_setting)
|
|
|
|
# e.g. bt0002
|
|
pr = PredictionResult(
|
|
[
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(=O)C(O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(O)C2=O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(=O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
|
|
],
|
|
0.17,
|
|
None,
|
|
)
|
|
|
|
mock_val = {"rule_triggered": True, "transformations": [pr]}
|
|
|
|
with patch.object(self.mock_setting, "expand", return_value=mock_val):
|
|
spw.predict_step(from_depth=0)
|
|
|
|
self.assertEqual(len(spw.smiles_to_node.keys()), 4)
|
|
self.assertEqual(len(spw.edges), 3)
|
|
|
|
json_result = spw.to_json()
|
|
|
|
self.assertIsInstance(json_result, dict)
|
|
self.assertIn("nodes", json_result)
|
|
self.assertIn("edges", json_result)
|
|
self.assertEqual(len(json_result["nodes"]), 4)
|
|
self.assertEqual(len(json_result["edges"]), 3)
|