sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2024-11-27 17:35:07 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Setup Search with Search_Config to Fix Tests

Browse source 
- Rename pytest fixture search_config to more appropriate
content_config
- Create search_config pytest fixture
- Use search_config where search being setup, used in tests
This commit is contained in:
Debanjum Singh Solanky 2022-01-14 20:13:14 -05:00
parent c64e0c2965
commit ed144f7984
4 changed files with 80 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
tests/conftest.py
View file

@ -1,51 +1,78 @@
# Standard Packages # Standard Packages
import pytest import pytest
from pathlib import Path from pathlib import Path
from src import search_type
# Internal Packages # Internal Packages
from src.search_type import asymmetric, image_search from src.search_type import asymmetric, image_search
from src.utils.rawconfig import ContentTypeConfig, ImageSearchConfig, TextSearchConfig from src.utils.rawconfig import AsymmetricConfig, ContentTypeConfig, ImageSearchConfig, ImageSearchTypeConfig, SearchTypeConfig, SymmetricConfig, TextSearchConfig
@pytest.fixture(scope='session') @pytest.fixture(scope='session')
def model_dir(tmp_path_factory): def search_config(tmp_path_factory):
model_dir = tmp_path_factory.mktemp('data') model_dir = tmp_path_factory.mktemp('data')
search_config = SearchTypeConfig()
search_config.asymmetric = SymmetricConfig(
encoder = "sentence-transformers/paraphrase-MiniLM-L6-v2",
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
model_directory = model_dir
)
search_config.asymmetric = AsymmetricConfig(
encoder = "sentence-transformers/msmarco-MiniLM-L-6-v3",
cross_encoder = "cross-encoder/ms-marco-MiniLM-L-6-v2",
model_directory = model_dir
)
search_config.image = ImageSearchTypeConfig(
encoder = "clip-ViT-B-32",
model_directory = model_dir
)
return search_config
@pytest.fixture(scope='session')
def model_dir(search_config):
model_dir = search_config.asymmetric.model_directory
# Generate Image Embeddings from Test Images # Generate Image Embeddings from Test Images
search_config = ContentTypeConfig() content_config = ContentTypeConfig()
search_config.image = ImageSearchConfig( content_config.image = ImageSearchConfig(
input_directory = 'tests/data', input_directory = 'tests/data',
embeddings_file = model_dir.joinpath('.image_embeddings.pt'), embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
batch_size = 10, batch_size = 10,
use_xmp_metadata = False) use_xmp_metadata = False)
image_search.setup(search_config.image, regenerate=False, verbose=True) image_search.setup(content_config.image, search_config.image, regenerate=False, verbose=True)
# Generate Notes Embeddings from Test Notes # Generate Notes Embeddings from Test Notes
search_config.org = TextSearchConfig( content_config.org = TextSearchConfig(
input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'], input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'],
input_filter = None, input_filter = None,
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'), compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
embeddings_file = model_dir.joinpath('.note_embeddings.pt')) embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
asymmetric.setup(search_config.org, regenerate=False, verbose=True) asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False, verbose=True)
return model_dir return model_dir
@pytest.fixture(scope='session') @pytest.fixture(scope='session')
def search_config(model_dir): def content_config(model_dir):
search_config = ContentTypeConfig() content_config = ContentTypeConfig()
search_config.org = TextSearchConfig( content_config.org = TextSearchConfig(
input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'], input_files = ['tests/data/main_readme.org', 'tests/data/interface_emacs_readme.org'],
input_filter = None, input_filter = None,
compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'), compressed_jsonl = model_dir.joinpath('.notes.jsonl.gz'),
embeddings_file = model_dir.joinpath('.note_embeddings.pt')) embeddings_file = model_dir.joinpath('.note_embeddings.pt'))
search_config.image = ImageSearchConfig( content_config.image = ImageSearchConfig(
input_directory = 'tests/data', input_directory = 'tests/data',
embeddings_file = 'tests/data/.image_embeddings.pt', embeddings_file = model_dir.joinpath('.image_embeddings.pt'),
batch_size = 10, batch_size = 10,
use_xmp_metadata = False) use_xmp_metadata = False)
return search_config return content_config

9
tests/test_asymmetric_search.py
View file

@ -1,14 +1,15 @@
# Internal Packages # Internal Packages
from src.main import model from src.main import model
from src.search_type import asymmetric from src.search_type import asymmetric
from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
# Test # Test
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_asymmetric_setup(search_config): def test_asymmetric_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Act # Act
# Regenerate notes embeddings during asymmetric setup # Regenerate notes embeddings during asymmetric setup
notes_model = asymmetric.setup(search_config.org, regenerate=True) notes_model = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=True)
# Assert # Assert
assert len(notes_model.entries) == 10 assert len(notes_model.entries) == 10
@ -16,9 +17,9 @@ def test_asymmetric_setup(search_config):
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_asymmetric_search(search_config): def test_asymmetric_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
model.notes_search = asymmetric.setup(search_config.org, regenerate=False) model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
query = "How to git install application?" query = "How to git install application?"
# Act # Act

51
tests/test_client.py
View file

@ -9,7 +9,7 @@ import pytest
from src.main import app, model, config from src.main import app, model, config
from src.search_type import asymmetric, image_search from src.search_type import asymmetric, image_search
from src.utils.helpers import resolve_absolute_path from src.utils.helpers import resolve_absolute_path
from src.utils.rawconfig import ContentTypeConfig from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
# Arrange # Arrange
@ -18,55 +18,60 @@ client = TestClient(app)
# Test # Test
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_search_with_invalid_search_type(): def test_search_with_invalid_content_type():
# Arrange # Arrange
user_query = "How to call semantic search from Emacs?" user_query = "How to call semantic search from Emacs?"
# Act # Act
response = client.get(f"/search?q={user_query}&t=invalid_search_type") response = client.get(f"/search?q={user_query}&t=invalid_content_type")
# Assert # Assert
assert response.status_code == 422 assert response.status_code == 422
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_search_with_valid_search_type(search_config: ContentTypeConfig): def test_search_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
config.content_type = search_config config.content_type = content_config
config.search_type = search_config
# config.content_type.image = search_config.image # config.content_type.image = search_config.image
for search_type in ["notes", "ledger", "music", "image"]: for content_type in ["notes", "ledger", "music", "image"]:
# Act # Act
response = client.get(f"/search?q=random&t={search_type}") response = client.get(f"/search?q=random&t={content_type}")
# Assert # Assert
assert response.status_code == 200 assert response.status_code == 200
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_regenerate_with_invalid_search_type(): def test_regenerate_with_invalid_content_type():
# Act # Act
response = client.get(f"/regenerate?t=invalid_search_type") response = client.get(f"/regenerate?t=invalid_content_type")
# Assert # Assert
assert response.status_code == 422 assert response.status_code == 422
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_regenerate_with_valid_search_type(search_config: ContentTypeConfig): def test_regenerate_with_valid_content_type(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
config.content_type = search_config config.content_type = content_config
for search_type in ["notes", "ledger", "music", "image"]: config.search_type = search_config
for content_type in ["notes", "ledger", "music", "image"]:
# Act # Act
response = client.get(f"/regenerate?t={search_type}") response = client.get(f"/regenerate?t={content_type}")
# Assert # Assert
assert response.status_code == 200 assert response.status_code == 200
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
@pytest.mark.skip(reason="Flaky test. Search doesn't always return expected image path.") @pytest.mark.skip(reason="Flaky test. Search doesn't always return expected image path.")
def test_image_search(search_config: ContentTypeConfig): def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
config.content_type = search_config config.content_type = content_config
model.image_search = image_search.setup(search_config.image, regenerate=False) config.search_type = search_config
model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
query_expected_image_pairs = [("brown kitten next to fallen plant", "kitten_park.jpg"), query_expected_image_pairs = [("brown kitten next to fallen plant", "kitten_park.jpg"),
("a horse and dog on a leash", "horse_dog.jpg"), ("a horse and dog on a leash", "horse_dog.jpg"),
("A guinea pig eating grass", "guineapig_grass.jpg")] ("A guinea pig eating grass", "guineapig_grass.jpg")]
@ -78,16 +83,16 @@ def test_image_search(search_config: ContentTypeConfig):
# Assert # Assert
assert response.status_code == 200 assert response.status_code == 200
actual_image = Path(response.json()[0]["Entry"]) actual_image = Path(response.json()[0]["Entry"])
expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name)) expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name))
# Assert # Assert
assert expected_image == actual_image assert expected_image == actual_image
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_notes_search(search_config: ContentTypeConfig): def test_notes_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
model.notes_search = asymmetric.setup(search_config.org, regenerate=False) model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
user_query = "How to git install application?" user_query = "How to git install application?"
# Act # Act
@ -101,9 +106,9 @@ def test_notes_search(search_config: ContentTypeConfig):
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_notes_search_with_include_filter(search_config: ContentTypeConfig): def test_notes_search_with_include_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
model.notes_search = asymmetric.setup(search_config.org, regenerate=False) model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
user_query = "How to git install application? +Emacs" user_query = "How to git install application? +Emacs"
# Act # Act
@ -117,9 +122,9 @@ def test_notes_search_with_include_filter(search_config: ContentTypeConfig):
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_notes_search_with_exclude_filter(search_config: ContentTypeConfig): def test_notes_search_with_exclude_filter(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
model.notes_search = asymmetric.setup(search_config.org, regenerate=False) model.notes_search = asymmetric.setup(content_config.org, search_config.asymmetric, regenerate=False)
user_query = "How to git install application? -clone" user_query = "How to git install application? -clone"
# Act # Act

13
tests/test_image_search.py
View file

@ -5,14 +5,15 @@ import pytest
from src.main import model from src.main import model
from src.search_type import image_search from src.search_type import image_search
from src.utils.helpers import resolve_absolute_path from src.utils.helpers import resolve_absolute_path
from src.utils.rawconfig import ContentTypeConfig, SearchTypeConfig
# Test # Test
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
def test_image_search_setup(search_config): def test_image_search_setup(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Act # Act
# Regenerate image search embeddings during image setup # Regenerate image search embeddings during image setup
image_search_model = image_search.setup(search_config.image, regenerate=True) image_search_model = image_search.setup(content_config.image, search_config.image, regenerate=True)
# Assert # Assert
assert len(image_search_model.image_names) == 3 assert len(image_search_model.image_names) == 3
@ -21,9 +22,9 @@ def test_image_search_setup(search_config):
# ---------------------------------------------------------------------------------------------------- # ----------------------------------------------------------------------------------------------------
@pytest.mark.skip(reason="results inconsistent currently") @pytest.mark.skip(reason="results inconsistent currently")
def test_image_search(search_config): def test_image_search(content_config: ContentTypeConfig, search_config: SearchTypeConfig):
# Arrange # Arrange
model.image_search = image_search.setup(search_config.image, regenerate=False) model.image_search = image_search.setup(content_config.image, search_config.image, regenerate=False)
query_expected_image_pairs = [("brown kitten next to plant", "kitten_park.jpg"), query_expected_image_pairs = [("brown kitten next to plant", "kitten_park.jpg"),
("horse and dog in a farm", "horse_dog.jpg"), ("horse and dog in a farm", "horse_dog.jpg"),
("A guinea pig eating grass", "guineapig_grass.jpg")] ("A guinea pig eating grass", "guineapig_grass.jpg")]
@ -38,11 +39,11 @@ def test_image_search(search_config):
results = image_search.collate_results( results = image_search.collate_results(
hits, hits,
model.image_search.image_names, model.image_search.image_names,
search_config.image.input_directory, content_config.image.input_directory,
count=1) count=1)
actual_image = results[0]["Entry"] actual_image = results[0]["Entry"]
expected_image = resolve_absolute_path(search_config.image.input_directory.joinpath(expected_image_name)) expected_image = resolve_absolute_path(content_config.image.input_directory.joinpath(expected_image_name))
# Assert # Assert
assert expected_image == actual_image assert expected_image == actual_image