Files
enviPy-bayer/tests/test_sobjects.py
2025-12-20 02:11:47 +13:00

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)