def test_request_1(self):
codex = ESCodex(multiplier=8)
request = Request()
request.body = REQUEST_BODY
request.url = URL(b'/test')
field, query = codex.parse_query(request)
self.assertEqual(field, 'passage')
self.assertEqual(query, b'this is a test')
topk, correct_cids = codex.multiply_request(request)
self.assertEqual(topk, 100)
size = load_json(request.body)['size']
self.assertEqual(size, 800)
def multiply_request(self, request: Request) -> Tuple[int, List[str]]:
"""Multiply size of Elasticsearch query"""
body = load_json(request.body)
topk = request.url.query.pop('size', None)
correct_cids = request.url.query.pop('nboost', None)
# search for topk in body
if body:
correct_cids = body.pop('nboost', correct_cids)
topk = body.pop('size', topk)
topk = 10 if topk is None else int(topk)
if body:
body['size'] = topk * self.multiplier
request.body = dump_json(body)
else:
request.url.query['size'] = str(topk * self.multiplier)
correct_cids = correct_cids.split(',') if correct_cids else None
return topk, correct_cids