chatqna-core: Fix backend unit test (open-edge-platform#212)

Signed-off-by: Yeoh, Hoong Tee <hoong.tee.yeoh@intel.com>

Author: hteeyeoh

sample-applications/chat-question-and-answer-core/app/chain.py

@@ -11,75 +11,85 @@
 from langchain_core.output_parsers import StrOutputParser
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.prompts import ChatPromptTemplate
+import os
 import pandas as pd
 
 config = Settings()
 vectorstore = None
 
-# login huggingface
-login_to_huggingface(config.HF_ACCESS_TOKEN)
-
-# Download convert the model to openvino optimized
-download_huggingface_model(config.EMBEDDING_MODEL_ID, config.CACHE_DIR)
-download_huggingface_model(config.RERANKER_MODEL_ID, config.CACHE_DIR)
-download_huggingface_model(config.LLM_MODEL_ID, config.CACHE_DIR)
-
-# Convert to openvino IR
-convert_model(config.EMBEDDING_MODEL_ID, config.CACHE_DIR, "embedding")
-convert_model(config.RERANKER_MODEL_ID, config.CACHE_DIR, "reranker")
-convert_model(config.LLM_MODEL_ID, config.CACHE_DIR, "llm")
-
-# Define RAG prompt
-template = """
-Use the following pieces of context from retrieved
-dataset to answer the question. Do not make up an answer if there is no
-context provided to help answer it.
-
-Context:
----------
-{context}
-
----------
-Question: {question}
----------
-
-Answer:
-"""
-
-prompt = ChatPromptTemplate.from_template(template)
-
-# Initialize Embedding Model
-embedding = OpenVINOBgeEmbeddings(
-    model_name_or_path=f"{config.CACHE_DIR}/{config.EMBEDDING_MODEL_ID}",
-    model_kwargs={"device": config.EMBEDDING_DEVICE, "compile": False},
-)
-embedding.ov_model.compile()
-
-# Initialize Reranker Model
-reranker = OpenVINOReranker(
-    model_name_or_path=f"{config.CACHE_DIR}/{config.RERANKER_MODEL_ID}",
-    model_kwargs={"device": config.RERANKER_DEVICE},
-    top_n=2,
-)
-
-# Initialize LLM
-llm = HuggingFacePipeline.from_model_id(
-    model_id=f"{config.CACHE_DIR}/{config.LLM_MODEL_ID}",
-    task="text-generation",
-    backend="openvino",
-    model_kwargs={
-        "device": config.LLM_DEVICE,
-        "ov_config": {
-            "PERFORMANCE_HINT": "LATENCY",
-            "NUM_STREAMS": "1",
-            "CACHE_DIR": f"{config.CACHE_DIR}/{config.LLM_MODEL_ID}/model_cache",
+# The RUN_TEST flag is used to bypass the model download and conversion steps during pytest unit testing.
+# By default, the flag is set to 'false', enabling the model download and conversion process in a normal run.
+# To skip these steps, set the flag to 'true'.
+# Check environment flag
+RUN_TEST = os.getenv('RUN_TEST', False)
+
+if not RUN_TEST:
+    # login huggingface
+    login_to_huggingface(config.HF_ACCESS_TOKEN)
+
+    # Download convert the model to openvino optimized
+    download_huggingface_model(config.EMBEDDING_MODEL_ID, config.CACHE_DIR)
+    download_huggingface_model(config.RERANKER_MODEL_ID, config.CACHE_DIR)
+    download_huggingface_model(config.LLM_MODEL_ID, config.CACHE_DIR)
+
+    # Convert to openvino IR
+    convert_model(config.EMBEDDING_MODEL_ID, config.CACHE_DIR, "embedding")
+    convert_model(config.RERANKER_MODEL_ID, config.CACHE_DIR, "reranker")
+    convert_model(config.LLM_MODEL_ID, config.CACHE_DIR, "llm")
+
+    # Define RAG prompt
+    template = """
+    Use the following pieces of context from retrieved
+    dataset to answer the question. Do not make up an answer if there is no
+    context provided to help answer it.
+
+    Context:
+    ---------
+    {context}
+
+    ---------
+    Question: {question}
+    ---------
+
+    Answer:
+    """
+
+    prompt = ChatPromptTemplate.from_template(template)
+
+    # Initialize Embedding Model
+    embedding = OpenVINOBgeEmbeddings(
+        model_name_or_path=f"{config.CACHE_DIR}/{config.EMBEDDING_MODEL_ID}",
+        model_kwargs={"device": config.EMBEDDING_DEVICE, "compile": False},
+    )
+    embedding.ov_model.compile()
+
+    # Initialize Reranker Model
+    reranker = OpenVINOReranker(
+        model_name_or_path=f"{config.CACHE_DIR}/{config.RERANKER_MODEL_ID}",
+        model_kwargs={"device": config.RERANKER_DEVICE},
+        top_n=2,
+    )
+
+    # Initialize LLM
+    llm = HuggingFacePipeline.from_model_id(
+        model_id=f"{config.CACHE_DIR}/{config.LLM_MODEL_ID}",
+        task="text-generation",
+        backend="openvino",
+        model_kwargs={
+            "device": config.LLM_DEVICE,
+            "ov_config": {
+                "PERFORMANCE_HINT": "LATENCY",
+                "NUM_STREAMS": "1",
+                "CACHE_DIR": f"{config.CACHE_DIR}/{config.LLM_MODEL_ID}/model_cache",
+            },
+            "trust_remote_code": True,
         },
-        "trust_remote_code": True,
-    },
-    pipeline_kwargs={"max_new_tokens": config.MAX_TOKENS},
-)
-if llm.pipeline.tokenizer.eos_token_id:
-    llm.pipeline.tokenizer.pad_token_id = llm.pipeline.tokenizer.eos_token_id
+        pipeline_kwargs={"max_new_tokens": config.MAX_TOKENS},
+    )
+    if llm.pipeline.tokenizer.eos_token_id:
+        llm.pipeline.tokenizer.pad_token_id = llm.pipeline.tokenizer.eos_token_id
+else:
+    logger.info("Bypassing to mock these functions because RUN_TEST is set to 'True' to run pytest unit test.")
 
 
 def default_context(docs):

sample-applications/chat-question-and-answer-core/tests/conftest.py

@@ -1,3 +1,9 @@
+import os
+
+# Configure the environment variable prior to importing the app
+# This ensures the app operates in test mode, bypassing the startup function responsible for model downloading and conversion
+os.environ['RUN_TEST'] = "True"
+
 import pytest
 from fastapi.testclient import TestClient
 

sample-applications/chat-question-and-answer-core/tests/test_server.py

@@ -3,9 +3,24 @@
 
 
 def test_chain_response(test_client, mocker):
+    """
+    Tests the chain response functionality of the server by simulating a POST
+    request to the `/stream_log` endpoint and verifying the streamed response.
+    Args:
+        test_client: A test client instance used to simulate HTTP requests.
+        mocker: A mocking library instance used to patch dependencies.
+    Mocks:
+        - `app.server.get_retriever`: Mocked to return `True`.
+        - `app.server.build_chain`: Mocked to return `True`.
+        - `app.server.process_query`: Mocked to return an iterator with values `["one", "two"]`.
+    Raises:
+        AssertionError: If any of the assertions fail.
+    """
 
     payload = {"input": "What is AI?", "stream": True}
 
+    mocker.patch("app.server.get_retriever", return_value=True)
+    mocker.patch("app.server.build_chain", return_value=True)
     mocker.patch("app.server.process_query", return_value=iter(["one", "two"]))
 
     response = test_client.post("/stream_log", json=payload)
@@ -22,6 +37,23 @@ def test_chain_response(test_client, mocker):
 
 
 def test_success_upload_and_create_embedding(test_client, mocker):
+    """
+    Tests the successful upload of a document and the creation of embeddings.
+    This test simulates the process of uploading a text file, validating the document,
+    saving it, and creating embeddings using a mocked FAISS vector database. It verifies
+    that the API endpoint responds with the correct status code and response JSON.
+    Args:
+        test_client: A test client instance used to simulate HTTP requests to the API.
+        mocker: A mocking library instance used to patch functions and simulate behavior.
+    Mocks:
+        - `app.server.validate_document`: Mocked to return `True`.
+        - `app.server.save_document`: Mocked to return the temporary file name and `None`.
+        - `app.server.create_faiss_vectordb`: Mocked to return `True`.
+    Assertions:
+        - The response status code is 200.
+        - The response JSON matches the expected success message and metadata.
+    """
+
     with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as tmp_file:
         tmp_file.write(b"This is sample txt file.")
         tmp_file.seek(0)
@@ -44,6 +76,21 @@ def test_success_upload_and_create_embedding(test_client, mocker):
 
 
 def test_success_get_documents(test_client, mocker):
+    """
+    Test the successful retrieval of documents from the server.
+    This test verifies that the `/documents` endpoint returns a 200 status code
+    and the expected JSON response containing a list of documents.
+    Args:
+        test_client (TestClient): A test client instance for making HTTP requests.
+        mocker (MockerFixture): A mocker fixture for patching and mocking dependencies.
+    Mocks:
+        - `app.server.get_document_from_vectordb`: Mocked to return a list of documents.
+    Assertions:
+        - The response status code is 200.
+        - The response JSON contains a "status" key with the value "Success".
+        - The response JSON contains a "metadata" key with a "documents" list matching the mocked documents.
+    """
+
     mock_documents = ["test1.txt", "test2.pdf"]
     mocker.patch('app.server.get_document_from_vectordb', return_value=mock_documents)
 
@@ -57,6 +104,20 @@ def test_success_get_documents(test_client, mocker):
 
 
 def test_delete_embedding_success(test_client, mocker):
+    """
+    Test the successful deletion of an embedding from the vector database.
+    This test verifies that the `delete_embedding_from_vectordb` function is called
+    and the API endpoint for deleting a document responds with the expected status code.
+    Args:
+        test_client: A test client instance for simulating HTTP requests to the server.
+        mocker: A mocking library instance used to patch and mock dependencies.
+    Mocks:
+        - `app.server.delete_embedding_from_vectordb`: Mocked to return `True`.
+    Assertions:
+        - Ensures that the HTTP DELETE request to the "/documents" endpoint with
+          the specified parameters returns a status code of HTTPStatus.NO_CONTENT.
+    """
+
     mocker.patch('app.server.delete_embedding_from_vectordb', return_value=True)
 
     response = test_client.delete("/documents", params={"document": "test1.txt"})
@@ -65,6 +126,20 @@ def test_delete_embedding_success(test_client, mocker):
 
 
 def test_delete_all_embedding_success(test_client, mocker):
+    """
+    Test the successful deletion of all embeddings from the vector database.
+    This test verifies that the endpoint for deleting all documents functions
+    correctly by mocking the `delete_embedding_from_vectordb` function to
+    return `True` and asserting that the response status code is `HTTPStatus.NO_CONTENT`.
+    Args:
+        test_client: A test client instance used to simulate HTTP requests to the server.
+        mocker: A mocking utility used to patch the `delete_embedding_from_vectordb` function.
+    Mocks:
+        - `app.server.delete_embedding_from_vectordb`: Mocked to return `True`.
+    Assertions:
+        - The response status code is `HTTPStatus.NO_CONTENT` (204).
+    """
+
     mocker.patch('app.server.delete_embedding_from_vectordb', return_value=True)
 
     response = test_client.delete("/documents", params={"delete_all": True})
@@ -73,6 +148,17 @@ def test_delete_all_embedding_success(test_client, mocker):
 
 
 def test_upload_unsupported_file(test_client):
+    """
+    Tests the upload of an unsupported file format to the server.
+    This test verifies that the server returns a 400 status code and an appropriate
+    error message when a file with an unsupported format (e.g., .html) is uploaded.
+    Args:
+        test_client: A test client instance used to simulate HTTP requests to the server.
+    Raises:
+        AssertionError: If the response status code is not 400 or the error message
+                        does not match the expected output.
+    """
+
     with tempfile.NamedTemporaryFile(delete=True, suffix=".html") as tmp_file:
         tmp_file.write(b"This is sample html file.")
         tmp_file.seek(0)
@@ -86,6 +172,21 @@ def test_upload_unsupported_file(test_client):
 
 
 def test_fail_get_documents(test_client, mocker):
+    """
+    Test case for handling failure when retrieving documents from the vector database.
+    This test simulates an exception being raised during the retrieval of documents
+    from the vector database and verifies that the server responds with the appropriate
+    HTTP status code and error message.
+    Args:
+        test_client: A test client instance used to simulate HTTP requests to the server.
+        mocker: A mocking library instance used to patch and simulate behavior of dependencies.
+    Mocks:
+        - `app.server.get_document_from_vectordb`: Mocked to raise an exception with the message "Error getting documents."
+    Asserts:
+        - The HTTP response status code is 500 (Internal Server Error).
+        - The JSON response contains the expected error message.
+    """
+
     mocker.patch('app.server.get_document_from_vectordb', side_effect=Exception("Error getting documents."))
 
     response = test_client.get("/documents")
@@ -97,6 +198,21 @@ def test_fail_get_documents(test_client, mocker):
 
 
 def test_delete_embedding_failure(test_client, mocker):
+    """
+    Test case for handling failure during the deletion of embeddings from the vector database.
+    This test simulates a failure scenario where the `delete_embedding_from_vectordb` function
+    raises an exception. It verifies that the server responds with the appropriate HTTP status
+    code and error message.
+    Args:
+        test_client (TestClient): A test client instance for simulating HTTP requests to the server.
+        mocker (MockerFixture): A fixture for mocking dependencies and functions.
+    Mocks:
+        - `app.server.delete_embedding_from_vectordb`: Mocked to raise an exception with the message "Error deleting embeddings."
+    Asserts:
+        - The response status code is 500 (Internal Server Error).
+        - The response JSON contains the expected error detail message.
+    """
+
     mocker.patch('app.server.delete_embedding_from_vectordb', side_effect=Exception("Error deleting embeddings."))
 
     response = test_client.delete("/documents", params={"document": "test1.txt"})