Files
enviPy-bayer/tests/test_dataset.py
jebus afeb56622c [Chore] Linted Files (#150)
Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#150
2025-10-09 07:25:13 +13:00

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 Dataset
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 = Dataset.generate_dataset(reactions, applicable_rules)
self.assertEqual(len(ds.y()), 1)
self.assertEqual(sum(ds.y()[0]), 1)