def test_services(storage_socket, session):
assert session.query(OptimizationProcedureORM).count() == 0
water = ptl.data.get_molecule("water_dimer_minima.psimol")
ret = storage_socket.add_molecules([water])
assert ret["meta"]["success"] is True
assert ret["meta"]["n_inserted"] == 1
proc_data = {
"initial_molecule": ret["data"][0],
"keywords": None,
"program": "p7",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
"protocols": {},
}
service_data = {
"tag": "tag1 tag2",
"hash_index": "123",
"status": "COMPLETE",
"optimization_program": "gaussian",
# extra fields
"torsiondrive_state": {},
"dihedral_template": "1", # Not realistic?
"optimization_template": "2", # Not realistic?
"molecule_template": "",
"logger": None,
"storage_socket": storage_socket,
"task_priority": 0,
}
procedure = OptimizationProcedureORM(**proc_data)
session.add(procedure)
session.commit()
assert procedure.id
service_pydantic = TorsionDriveService(**service_data)
doc = ServiceQueueORM(**service_pydantic.dict(
include=set(ServiceQueueORM.__dict__.keys())))
doc.extra = service_pydantic.dict(
exclude=set(ServiceQueueORM.__dict__.keys()))
doc.procedure_id = procedure.id
doc.priority = doc.priority.value # Special case where we need the value not the enum
session.add(doc)
session.commit()
session.delete(doc)
session.delete(procedure)
session.commit()
assert session.query(ServiceQueueORM).count() == 0
def test_services(storage_socket, session):
assert session.query(OptimizationProcedureORM).count() == 0
proc_data = {
"initial_molecule": None,
"keywords": None,
"program": "p7",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
}
service_data = {
"tag": "tag1 tag2",
"hash_index": "123",
"status": "COMPLETE",
"optimization_program": "gaussian",
# extra fields
"torsiondrive_state": {},
"dihedral_template": "1",
"optimization_template": "2",
"molecule_template": "",
"logger": None,
"storage_socket": storage_socket,
"task_priority": 0
}
procedure = OptimizationProcedureORM(**proc_data)
session.add(procedure)
session.commit()
assert procedure.id
service_pydantic = TorsionDriveService(**service_data)
doc = ServiceQueueORM(**service_pydantic.dict(
include=set(ServiceQueueORM.__dict__.keys())))
doc.extra = service_pydantic.dict(
exclude=set(ServiceQueueORM.__dict__.keys()))
doc.procedure_id = procedure.id
doc.priority = doc.priority.value # Special case where we need the value not the enum
session.add(doc)
session.commit()
session.delete(doc)
session.delete(procedure)
session.commit()
assert session.query(ServiceQueueORM).count() == 0
def test_services_sql(storage_results):
mol_ids = [int(mol.id) for mol in storage_results.get_molecules()['data']]
torsion_proc = {
"procedure": "torsiondrive",
"keywords": {
"dihedrals": [[0, 1, 2, 3]],
"grid_spacing": [10]
},
"hash_index": 456,
"optimization_spec": {
"program": "geometric",
"keywords": {
"coordsys": "tric",
}
},
"qc_spec": {
"driver": "gradient",
"method": "HF",
"basis": "sto-3g",
# "keywords": None,
"program": "psi4"
},
"initial_molecule": [mol_ids[0], mol_ids[1]],
"final_energy_dict": {},
"optimization_history": {},
"minimum_positions": {},
"provenance": {
"creator": ""
}
}
# Procedure
proc_pydantic = ptl.models.TorsionDriveRecord(**torsion_proc)
service_data = {
"tag": "tag1 tag2",
"hash_index": "123",
"status": TaskStatusEnum.waiting,
# extra fields
"torsiondrive_state": {},
"dihedral_template": "1",
"optimization_template": "2",
"molecule_template": "",
"logger": None,
"storage_socket": storage_results,
"task_priority": 0,
"output": proc_pydantic
}
service = TorsionDriveService(**service_data)
ret = storage_results.add_services([service])
assert len(ret['data']) == 1
ret = storage_results.get_services(procedure_id=ret['data'][0],
status=TaskStatusEnum.waiting)
assert ret['data'][0]['hash_index'] == service_data['hash_index']
# attributes in extra fields
assert ret['data'][0]['dihedral_template'] == service_data[
'dihedral_template']
# Create Pydantic object from DB returned object
py_obj = TorsionDriveService(**ret['data'][0],
storage_socket=storage_results,
logger=None)
assert py_obj
# Test update
py_obj.task_priority = 3
ret_count = storage_results.update_services([py_obj])
assert ret_count == 1
ret = storage_results.get_services(
procedure_id=ret['data'][0]['procedure_id'],
status=TaskStatusEnum.waiting)
assert ret['data'][0]['task_priority'] == py_obj.task_priority