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