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"]
def test_results_pagination(storage_socket):
"""
Test results pagination
"""
# results = storage_socket.get_results()['data']
# storage_socket.del_results([result['id'] for result in results])
assert len(storage_socket.get_results()['data']) == 0
water = ptl.data.get_molecule("water_dimer_minima.psimol")
mol = storage_socket.add_molecules([water])['data'][0]
result_template = {
"molecule": mol,
"method": "M1",
"basis": "B1",
"keywords": None,
"program": "P1",
"driver": "energy",
}
# Save (~ 1-7 msec/doc)
t1 = time()
total_results = 50
first_half = int(total_results / 2)
limit = 10
skip = 5
results = []
for i in range(first_half):
tmp = result_template.copy()
tmp['basis'] = str(i)
results.append(ptl.models.ResultRecord(**tmp))
result_template['method'] = 'M2'
for i in range(first_half, total_results):
tmp = result_template.copy()
tmp['basis'] = str(i)
results.append(ptl.models.ResultRecord(**tmp))
inserted = storage_socket.add_results(results)
assert inserted['meta']['n_inserted'] == total_results
# total_time = (time() - t1) * 1000 / total_results
# print('Inserted {} results in {:.2f} msec / doc'.format(total_results, total_time))
#
# query (~ 0.03 msec/doc)
# t1 = time()
ret = storage_socket.get_results(method='M2',
status=None,
limit=limit,
skip=skip)
# total_time = (time() - t1) * 1000 / first_half
# print('Query {} results in {:.2f} msec /doc'.format(first_half, total_time))
# count is total, but actual data size is the limit
assert ret['meta']['n_found'] == total_results - first_half
assert len(ret['data']) == limit
assert int(ret['data'][0]['basis']) == first_half + skip
# get the last page when with fewer than limit are remaining
ret = storage_socket.get_results(method='M1',
skip=(int(first_half - limit / 2)),
status=None)
assert len(ret['data']) == limit / 2
# cleanup
storage_socket.del_results(inserted['data'])
storage_socket.del_molecules(mol)
def test_results_add(storage_socket):
# Add two waters
water = ptl.data.get_molecule("water_dimer_minima.psimol")
water2 = ptl.data.get_molecule("water_dimer_stretch.psimol")
mol_insert = storage_socket.add_molecules([water, water2])
kw1 = ptl.models.KeywordSet(**{"comments": "a", "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",
# "extras": {
# "other_data": 5
# },
"hash_index": 0,
})
page2 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][1],
"method": "M1",
"basis": "B1",
"keywords": kwid1,
"program": "P1",
"driver": "energy",
# "extras": {
# "other_data": 10
# },
"hash_index": 1,
})
page3 = ptl.models.ResultRecord(
**{
"molecule": mol_insert["data"][1],
"method": "M22",
"basis": "B1",
"keywords": None,
"program": "P1",
"driver": "energy",
# "extras": {
# "other_data": 10
# },
"hash_index": 2,
})
ids = []
ret = storage_socket.add_results([page1, page2])
assert ret["meta"]["n_inserted"] == 2
ids.extend(ret['data'])
# add with duplicates:
ret = storage_socket.add_results([page1, page2, page3])
assert ret["meta"]["n_inserted"] == 1
assert len(ret['data']) == 3 # first 2 found are None
assert len(ret["meta"]['duplicates']) == 2
for res_id in ret['data']:
if res_id is not None:
ids.append(res_id)
ret = storage_socket.del_results(ids)
assert ret == 3
ret = storage_socket.del_molecules(id=mol_insert["data"])
assert ret == 2
