def test_vectorsearch_driver_mock_queryset():
# no queryset
driver = VectorSearchDriver(top_k=3)
assert driver.top_k == 3
# with queryset
driver = SimpleVectorSearchDriver(top_k=3)
assert driver.top_k == 4
def test_vectorsearch_driver_mock_queryset(self):
# no queryset
driver = VectorSearchDriver(top_k=3)
self.assertEqual(driver.top_k, 3)
# with queryset
driver = SimpleVectorSearchDriver(top_k=3)
self.assertEqual(driver.top_k, 4)
def search(*args, **kwargs):
"""Generate a searching request """
if ('top_k' in kwargs) and (kwargs['top_k'] is not None):
from jina.drivers.search import VectorSearchDriver
topk_ql = QueryLang(VectorSearchDriver(top_k=kwargs['top_k']))
if 'queryset' not in kwargs:
kwargs['queryset'] = [topk_ql]
else:
kwargs['queryset'].append(topk_ql)
yield from _generate(*args, **kwargs)
def test_topk_override(config):
# Making queryset
top_k_queryset = QueryLang(
VectorSearchDriver(top_k=int(os.environ['JINA_TOPK_OVERRIDE']),
priority=1))
with Flow.load_config('flow.yml') as index_flow:
index_flow.index(input_fn=random_docs(100))
with Flow.load_config('flow.yml') as search_flow:
search_flow.search(input_fn=random_docs(int(os.environ['JINA_NDOCS'])),
output_fn=validate_override_results,
queryset=[top_k_queryset])
def test_topk_override(config, mocker):
NDOCS = 3
TOPK_OVERRIDE = 11
def validate(resp):
mock()
assert len(resp.search.docs) == NDOCS
for doc in resp.search.docs:
assert len(doc.matches) == TOPK_OVERRIDE
# Making queryset
top_k_queryset = QueryLang(
VectorSearchDriver(top_k=TOPK_OVERRIDE, priority=1))
with Flow.load_config('flow.yml') as index_flow:
index_flow.index(input_fn=random_docs(100))
mock = mocker.Mock()
with Flow.load_config('flow.yml') as search_flow:
search_flow.search(input_fn=random_docs(NDOCS),
on_done=validate,
queryset=[top_k_queryset])
mock.assert_called_once()