sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2025-02-17 08:04:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use async/await in tests for query method of text and image search

Browse source 
The text, image search query method has become async. So async/await
is required to get results correctly in tests etc
This commit is contained in:
Debanjum Singh Solanky 2023-06-28 20:11:26 -07:00
parent f516d127c8
commit 56ce97ef9e
3 changed files with 16 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/khoj/routers/api.py
View file

@ -157,7 +157,7 @@ async def search(
user_query = q.strip()
results_count = n
score_threshold = score_threshold if score_threshold is not None else -math.inf
search_futures: list[concurrent.futures.Future] = []
search_futures: List[concurrent.futures.Future] = []
# return cached results, if available
query_cache_key = f"{user_query}-{n}-{t}-{r}-{score_threshold}-{dedupe}"

18
tests/test_image_search.py
View file

@ -3,6 +3,9 @@ import logging
from pathlib import Path
from PIL import Image
# External Packages
import pytest
# Internal Packages
from khoj.utils.state import model
from khoj.utils.constants import web_directory
@ -48,7 +51,8 @@ def test_image_metadata(content_config: ContentConfig):
# ----------------------------------------------------------------------------------------------------
def test_image_search(content_config: ContentConfig, search_config: SearchConfig):
@pytest.mark.anyio
async def test_image_search(content_config: ContentConfig, search_config: SearchConfig):
# Arrange
output_directory = resolve_absolute_path(web_directory)
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
@ -60,7 +64,7 @@ def test_image_search(content_config: ContentConfig, search_config: SearchConfig
# Act
for query, expected_image_name in query_expected_image_pairs:
hits = image_search.query(query, count=1, model=model.image_search)
hits = await image_search.query(query, count=1, model=model.image_search)
results = image_search.collate_results(
hits,
@ -83,7 +87,8 @@ def test_image_search(content_config: ContentConfig, search_config: SearchConfig
# ----------------------------------------------------------------------------------------------------
def test_image_search_query_truncated(content_config: ContentConfig, search_config: SearchConfig, caplog):
@pytest.mark.anyio
async def test_image_search_query_truncated(content_config: ContentConfig, search_config: SearchConfig, caplog):
# Arrange
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
max_words_supported = 10
@ -93,7 +98,7 @@ def test_image_search_query_truncated(content_config: ContentConfig, search_conf
# Act
try:
with caplog.at_level(logging.INFO, logger="khoj.search_type.image_search"):
image_search.query(query, count=1, model=model.image_search)
await image_search.query(query, count=1, model=model.image_search)
# Assert
except RuntimeError as e:
if "The size of tensor a (102) must match the size of tensor b (77)" in str(e):
@ -102,7 +107,8 @@ def test_image_search_query_truncated(content_config: ContentConfig, search_conf
# ----------------------------------------------------------------------------------------------------
def test_image_search_by_filepath(content_config: ContentConfig, search_config: SearchConfig, caplog):
@pytest.mark.anyio
async def test_image_search_by_filepath(content_config: ContentConfig, search_config: SearchConfig, caplog):
# Arrange
output_directory = resolve_absolute_path(web_directory)
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
@ -113,7 +119,7 @@ def test_image_search_by_filepath(content_config: ContentConfig, search_config:
# Act
with caplog.at_level(logging.INFO, logger="khoj.search_type.image_search"):
hits = image_search.query(query, count=1, model=model.image_search)
hits = await image_search.query(query, count=1, model=model.image_search)
results = image_search.collate_results(
hits,

5
tests/test_text_search.py
View file

@ -72,13 +72,14 @@ def test_text_content_index_only_updates_on_changes(content_config: ContentConfi
# ----------------------------------------------------------------------------------------------------
def test_asymmetric_search(content_config: ContentConfig, search_config: SearchConfig):
@pytest.mark.anyio
async def test_asymmetric_search(content_config: ContentConfig, search_config: SearchConfig):
# Arrange
model.notes_search = text_search.setup(OrgToJsonl, content_config.org, search_config.asymmetric, regenerate=True)
query = "How to git install application?"
# Act
hits, entries = text_search.query(query, model=model.notes_search, rank_results=True)
hits, entries = await text_search.query(query, model=model.notes_search, rank_results=True)
results = text_search.collate_results(hits, entries, count=1)