def test_mol_pagination(storage_socket):
"""
Test Molecule pagination
"""
assert len(storage_socket.get_molecules()['data']) == 0
mol_names = [
'water_dimer_minima.psimol', 'water_dimer_stretch.psimol',
'water_dimer_stretch2.psimol', 'neon_tetramer.psimol'
]
total = len(mol_names)
molecules = []
for mol_name in mol_names:
mol = ptl.data.get_molecule(mol_name)
molecules.append(mol)
inserted = storage_socket.add_molecules(molecules)
assert inserted['meta']['n_inserted'] == total
ret = storage_socket.get_molecules(skip=1)
assert len(ret['data']) == total - 1
assert ret['meta']['n_found'] == total
ret = storage_socket.get_molecules(skip=total + 1)
assert len(ret['data']) == 0
assert ret['meta']['n_found'] == total
# cleanup
storage_socket.del_molecules(inserted['data'])
def test_mol_formula(storage_socket):
"""
Test Molecule pagination
"""
assert len(storage_socket.get_molecules()["data"]) == 0
mol_names = [
"water_dimer_minima.psimol",
]
total = len(mol_names)
molecules = []
for mol_name in mol_names:
mol = ptl.data.get_molecule(mol_name)
molecules.append(mol)
inserted = storage_socket.add_molecules(molecules)
try:
assert inserted["meta"]["n_inserted"] == total
ret = storage_socket.get_molecules(molecular_formula="H4O2")
assert len(ret["data"]) == 1
assert ret["meta"]["n_found"] == 1
ret = storage_socket.get_molecules(molecular_formula="O2H4")
assert len(ret["data"]) == 1
assert ret["meta"]["n_found"] == 1
ret = storage_socket.get_molecules(molecular_formula="H4o2")
assert len(ret["data"]) == 0
assert ret["meta"]["n_found"] == 0
finally:
# cleanup
storage_socket.del_molecules(inserted["data"])
def test_mol_pagination(storage_socket):
"""
Test Molecule pagination
"""
assert len(storage_socket.get_molecules()["data"]) == 0
mol_names = [
"water_dimer_minima.psimol",
"water_dimer_stretch.psimol",
"water_dimer_stretch2.psimol",
"neon_tetramer.psimol",
]
total = len(mol_names)
molecules = []
for mol_name in mol_names:
mol = ptl.data.get_molecule(mol_name)
molecules.append(mol)
inserted = storage_socket.add_molecules(molecules)
try:
assert inserted["meta"]["n_inserted"] == total
ret = storage_socket.get_molecules(skip=1)
assert len(ret["data"]) == total - 1
assert ret["meta"]["n_found"] == total
ret = storage_socket.get_molecules(skip=total + 1)
assert len(ret["data"]) == 0
assert ret["meta"]["n_found"] == total
finally:
# cleanup
storage_socket.del_molecules(inserted["data"])
def test_molecules_add(storage_socket):
water = ptl.data.get_molecule("water_dimer_minima.psimol")
# Add once
ret1 = storage_socket.add_molecules([water])
assert ret1["meta"]["success"] is True
assert ret1["meta"]["n_inserted"] == 1
# Try duplicate adds
ret2 = storage_socket.add_molecules([water])
assert ret2["meta"]["success"] is True
assert ret2["meta"]["n_inserted"] == 0
assert ret2["meta"]["duplicates"][0] == ret1["data"][0]
# Assert the ids match
assert ret1["data"][0] == ret2["data"][0]
# Pull molecule from the DB for tests
db_json = storage_socket.get_molecules(
molecule_hash=water.get_hash())["data"][0]
water.compare(db_json)
# Cleanup adds
ret = storage_socket.del_molecules(molecule_hash=water.get_hash())
assert ret == 1
def test_molecule_sql(storage_socket, session):
"""
Test the use of the ME class MoleculeORM
Note:
creation of a MoleculeORM using ME is not implemented yet
Should create a MoleculeORM using: mongoengine_socket.add_molecules
"""
num_mol_in_db = session.query(MoleculeORM).count()
# MoleculeORM.objects().delete()
assert num_mol_in_db == 0
water = ptl.data.get_molecule("water_dimer_minima.psimol")
water2 = ptl.data.get_molecule("water_dimer_stretch.psimol")
# Add MoleculeORM
ret = storage_socket.add_molecules([water, water2])
assert ret["meta"]["success"] is True
assert ret["meta"]["n_inserted"] == 2
ret = storage_socket.get_molecules()
assert ret['meta']['n_found'] == 2
# Use the ORM class
water_mol = session.query(MoleculeORM).first()
assert water_mol.molecular_formula == "H4O2"
assert water_mol.molecular_charge == 0
# print(water_mol.dict())
#
# Query with fields in the model
result_list = session.query(MoleculeORM).filter_by(
molecular_formula="H4O2").all()
assert len(result_list) == 2
assert result_list[0].molecular_multiplicity == 1
# Query with fields NOT in the model. works too!
result_list = session.query(MoleculeORM).filter_by(
molecular_charge=0).all()
assert len(result_list) == 2
# get unique by hash and formula
one_mol = session.query(MoleculeORM).filter_by(
molecule_hash=water_mol.molecule_hash,
molecular_formula=water_mol.molecular_formula)
assert len(one_mol.all()) == 1
# Clean up
storage_socket.del_molecules(
molecule_hash=[water.get_hash(), water2.get_hash()])
def test_molecules_get(storage_socket):
water = ptl.data.get_molecule("water_dimer_minima.psimol")
# Add once
ret = storage_socket.add_molecules([water])
assert ret["meta"]["n_inserted"] == 1
water_id = ret["data"][0]
# Pull molecule from the DB for tests
water2 = storage_socket.get_molecules(id=water_id)["data"][0]
water2.compare(water)
# Cleanup adds
ret = storage_socket.del_molecules(id=water_id)
assert ret == 1
def test_molecules_bad_get(storage_socket):
water = ptl.data.get_molecule("water_dimer_minima.psimol")
# Add once
ret = storage_socket.add_molecules([water])
water_id = ret["data"][0]
# Pull molecule from the DB for tests
ret = storage_socket.get_molecules(id=[water_id, bad_id1, bad_id2])
assert ret["data"][0].id == water_id
assert ret["meta"]["n_found"] == 1
# Cleanup adds
ret = storage_socket.del_molecules(id=water_id)
assert ret == 1
def test_molecules_bad_get(storage_socket):
water = ptl.data.get_molecule("water_dimer_minima.psimol")
# Add once
ret = storage_socket.add_molecules([water])
water_id = ret["data"][0]
# Pull molecule from the DB for tests
ret = storage_socket.get_molecules(id=[water_id, "something", 5, (3, 2)])
assert len(ret["meta"]["errors"]) == 1
assert ret["meta"]["errors"][0][0] == "id"
assert len(ret["meta"]["errors"][0][1]) == 3
assert ret["meta"]["n_found"] == 1
# Cleanup adds
ret = storage_socket.del_molecules(id=water_id)
assert ret == 1
def test_identical_mol_insert(storage_socket):
"""
Tests as edge case where to identical molecules are added under different tags.
"""
water = ptl.data.get_molecule("water_dimer_minima.psimol")
# Add two identical molecules
ret1 = storage_socket.add_molecules([water, water])
assert ret1["meta"]["success"] is True
assert ret1["meta"]["n_inserted"] == 1
assert ret1["data"][0] == ret1["data"][1]
# Should only find one molecule
ret2 = storage_socket.get_molecules(molecule_hash=[water.get_hash()])
assert ret2["meta"]["n_found"] == 1
ret = storage_socket.del_molecules(molecule_hash=water.get_hash())
assert ret == 1
def storage_results(storage_socket):
# Add two waters
assert len(storage_socket.get_molecules()['data']) == 0
mol_names = [
'water_dimer_minima.psimol', 'water_dimer_stretch.psimol',
'water_dimer_stretch2.psimol', 'neon_tetramer.psimol'
]
molecules = []
for mol_name in mol_names:
mol = ptl.data.get_molecule(mol_name)
molecules.append(mol)
mol_insert = storage_socket.add_molecules(molecules)
kw1 = ptl.models.KeywordSet(**{"values": {}})
kwid1 = storage_socket.add_keywords([kw1])["data"][0]
page1 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][0],
"method": "M1",
"basis": "B1",
"keywords": kwid1,
"program": "P1",
"driver": "energy",
"return_result": 5,
"hash_index": 0,
"status": 'COMPLETE'
})
page2 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][1],
"method": "M1",
"basis": "B1",
"keywords": kwid1,
"program": "P1",
"driver": "energy",
"return_result": 10,
"hash_index": 1,
"status": 'COMPLETE'
})
page3 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][0],
"method": "M1",
"basis": "B1",
"keywords": kwid1,
"program": "P2",
"driver": "gradient",
"return_result": 15,
"hash_index": 2,
"status": 'COMPLETE'
})
page4 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][0],
"method": "M2",
"basis": "B1",
"keywords": kwid1,
"program": "P2",
"driver": "gradient",
"return_result": 15,
"hash_index": 3,
"status": 'COMPLETE'
})
page5 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][1],
"method": "M2",
"basis": "B1",
"keywords": kwid1,
"program": "P1",
"driver": "gradient",
"return_result": 20,
"hash_index": 4,
"status": 'COMPLETE'
})
page6 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][1],
"method": "M3",
"basis": "B1",
"keywords": None,
"program": "P1",
"driver": "gradient",
"return_result": 20,
"hash_index": 5,
"status": 'COMPLETE'
})
results_insert = storage_socket.add_results(
[page1, page2, page3, page4, page5, page6])
assert results_insert["meta"]["n_inserted"] == 6
yield storage_socket
# Cleanup
all_tasks = storage_socket.get_queue()['data']
storage_socket.del_tasks(id=[task.id for task in all_tasks])
result_ids = [x for x in results_insert["data"]]
ret = storage_socket.del_results(result_ids)
assert ret == results_insert["meta"]["n_inserted"]
ret = storage_socket.del_molecules(id=mol_insert["data"])
assert ret == mol_insert["meta"]["n_inserted"]
