Coverage for tests/test_query_database.py: 100%

1import pytest 

2from flask import current_app 

3from app.doc_parsers.parse_pdf import parse_pdf 

4from app.doc_indexer.index_doc import index_and_add_to_db 

5from app.doc_indexer.retrieve_document import query_database 

6from app import db 

7from sqlalchemy import delete 

8 

9 

10def test_query_vector_db_with_doc(client, app, clean_vector_db): 

11 """ 

12 Tests to make sure query database fuction returns a document when document exists 

13 """ 

14 FILE_PATH = "./tests/test_data/Dna.pdf" 

15 

16 # Parse the PDF into document chunks 

17 chunks = parse_pdf(FILE_PATH) 

18 assert len(chunks) > 0, "Parsing failed, no chunks extracted" 

19 

20 with app.app_context(): 

21 # Index chunks into the vector database 

22 index_and_add_to_db(chunks) 

23 vector_db = current_app.vector_db 

24 indexed_chunks = db.session.query(vector_db.EmbeddingStore.id).all() 

25 

26 assert len(indexed_chunks) == len( 

27 chunks 

28 ), "Not all chunks were indexed successfully" 

29 

30 with app.app_context(): 

31 documents = query_database("dna") 

32 

33 assert documents 

34 doc, score = documents[0] 

35 assert doc 

36 assert score >= 0 

37 assert score <= 1 

38 

39 

40def test_query_vector_db_no_doc(client, app, clean_vector_db): 

41 """ 

42 Tests to make sure query database fuction returns a does not return a document when one doesnt exist 

43 """ 

44 

45 with app.app_context(): 

46 documents = query_database("dna") 

47 

48 assert not documents 

49 

50 

51def test_embedding_model(app): 

52 with app.app_context(): 

53 assert current_app.vector_db.embeddings.model == "mxbai-embed-large"