forked from enviPath/enviPy
53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
from django.test import TestCase
|
|
|
|
from epdb.logic import PackageManager
|
|
from epdb.models import Reaction, Compound, User, Rule
|
|
from utilities.ml import RuleBasedDataset
|
|
|
|
|
|
class DatasetTest(TestCase):
|
|
fixtures = ["test_fixtures.jsonl.gz"]
|
|
|
|
def setUp(self):
|
|
self.cs1 = Compound.create(
|
|
self.package,
|
|
name="2,6-Dibromohydroquinone",
|
|
description="http://localhost:8000/package/32de3cf4-e3e6-4168-956e-32fa5ddb0ce1/compound/d6435251-1a54-4327-b4b1-fd6e9a8f4dc9/structure/d8a0225c-dbb5-4e6c-a642-730081c09c5b",
|
|
smiles="C1=C(C(=C(C=C1O)Br)O)Br",
|
|
).default_structure
|
|
|
|
self.cs2 = Compound.create(
|
|
self.package,
|
|
smiles="O=C(O)CC(=O)/C=C(/Br)C(=O)O",
|
|
).default_structure
|
|
|
|
self.rule1 = Rule.create(
|
|
rule_type="SimpleAmbitRule",
|
|
package=self.package,
|
|
smirks="[#8:8]([H])-[c:4]1[c:3]([H])[c:2](-[#1,#17,#35:9])[c:1](-[#8:7]([H]))[c:6](-[#1,#17,#35])[c:5]([H])1>>[#8-]-[#6:6](=O)-[#6:5]-[#6:4](=[O:8])\\[#6:3]=[#6:2](\\[#1,#17,#35:9])-[#6:1](-[#8-])=[O:7]",
|
|
description="http://localhost:8000/package/32de3cf4-e3e6-4168-956e-32fa5ddb0ce1/simple-ambit-rule/f6a56c0f-a4a0-4ee3-b006-d765b4767cf6",
|
|
)
|
|
|
|
self.reaction1 = Reaction.create(
|
|
package=self.package,
|
|
educts=[self.cs1],
|
|
products=[self.cs2],
|
|
rules=[self.rule1],
|
|
multi_step=False,
|
|
)
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
super(DatasetTest, cls).setUpClass()
|
|
cls.user = User.objects.get(username="anonymous")
|
|
cls.package = PackageManager.create_package(cls.user, "Anon Test Package", "No Desc")
|
|
|
|
def test_smoke(self):
|
|
reactions = [r for r in Reaction.objects.filter(package=self.package)]
|
|
applicable_rules = [self.rule1]
|
|
|
|
ds = RuleBasedDataset.generate_dataset(reactions, applicable_rules)
|
|
|
|
self.assertEqual(len(ds.y()), 1)
|
|
self.assertEqual(sum(ds.y()[0]), 1)
|