Merge pull request #103 from AET-DevOps25/feature/add-genai-tests

Add unit and integration tests for genai module

Author: edkaya

.github/workflows/genai-tests.yml

@@ -0,0 +1,46 @@
+name: CI/CD Tests in GenAI
+
+on:
+  push:
+    branches-ignore:
+      - main
+    paths:
+      - 'genai/**'
+      - '.github/workflows/genai-tests.yml'
+
+  pull_request:
+    branches-ignore:
+      - main
+    paths:
+      - 'genai/**'
+      - '.github/workflows/genai-tests.yml'
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    # Setup qdrant for integration tests
+    services:
+      qdrant:
+        image: qdrant/qdrant
+        ports:
+          - 6333:6333
+    env:
+      OPENAI_API_KEY: ${{ secrets.API_OPENAI }}
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Install dependencies
+        run: pip install -r requirements.txt
+        working-directory: genai
+
+      - name: Run genai tests
+        run: |
+          cd genai
+          pytest

genai/pytest.ini

@@ -0,0 +1,3 @@
+# pytest.ini
+[pytest]
+pythonpath = .

genai/requirements.txt

@@ -26,4 +26,9 @@ pydantic==2.11.5
 uvicorn[standard]==0.34.3
 
 # File Upload
-python-multipart==0.0.20
+python-multipart==0.0.20
+
+# Testing
+pytest==8.4.1
+fpdf==1.7.2
+pypdf==5.6.0

genai/service/rag_service.py

@@ -35,6 +35,8 @@ def prepare_prompt(system_prompt: str,
 
 def process_raw_messages(raw_messages: List[Dict]) -> List[BaseMessage]:
     """Turns raw messages into BaseMessages, so they can be passed into LLM"""
+    if not raw_messages:
+        return []
     processed_messages = []
     for msg in raw_messages:
         role = msg.get("role")

genai/tests/integration/generation_test.py

@@ -0,0 +1,62 @@
+from fastapi.testclient import TestClient
+from unittest.mock import patch, MagicMock
+from main import app
+
+client = TestClient(app)
+
+
+@patch("rag.llm.chat_model.ChatModel.invoke")
+@patch("routes.routes.qdrant.client.collection_exists", return_value=False)
+def test_generate_endpoint_success(_mock_exists, mock_invoke):
+    mock_response = MagicMock()
+    mock_response.content = "This is a test response"
+    mock_invoke.return_value = mock_response
+
+    payload = {
+        "query": "What can I cook with rice?",
+        "messages": [
+            {"role": "USER", "content": "I have rice and eggs"},
+            {"role": "ASSISTANT", "content": "How about fried rice?"}
+        ]
+    }
+
+    response = client.post("/genai/generate", json=payload)
+
+    assert response.status_code == 200
+    data = response.json()
+    assert "response" in data
+    assert data["response"] == "This is a test response"
+    mock_invoke.assert_called_once()
+
+
+@patch("rag.llm.chat_model.ChatModel.invoke")
+@patch("routes.routes.qdrant.client.collection_exists", return_value=False)
+def test_generate_endpoint_empty_messages(_mock_exists, mock_invoke):
+
+    mock_response = MagicMock()
+    mock_response.content = "No prior messages, here's a new idea!"
+    mock_invoke.return_value = mock_response
+
+    payload = {
+        "query": "Can I cook with lentils?",
+        "messages": []
+    }
+
+    response = client.post("/genai/generate", json=payload)
+
+    assert response.status_code == 200
+    data = response.json()
+    assert "response" in data
+    assert data["response"] == "No prior messages, here's a new idea!"
+
+
+def test_generate_endpoint_missing_fields():
+    payload = {
+        "query": "Can I cook with lentils?"
+        # "messages" key is missing
+    }
+
+    response = client.post("/genai/generate", json=payload)
+
+    assert response.status_code == 400
+    assert response.json() == {"detail": "Missing 'query' or 'messages'"}

genai/tests/integration/ingestion_test.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+from fpdf import FPDF
+from rag.ingestion_pipeline import IngestionPipeline
+from vector_database.qdrant_vdb import QdrantVDB
+
+
+# Helper method to generate a dummy pdf file with real content
+def generate_sample_pdf(path: Path, text: str = "Real ingestion test."):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, text)
+    pdf.output(str(path))
+
+
+def test_real_ingestion_pipeline(tmp_path):
+    collection_name = "test_collection"
+    qdrant = QdrantVDB()
+    # Just for testing purposes
+    qdrant.host = "http://localhost:6333"
+    qdrant.client = qdrant.get_vector_database(qdrant.host)
+
+    # Ensure collection does not exists before tests
+    qdrant.delete_collection(collection_name)
+
+    # Create a dummy PDF
+    pdf_path = tmp_path / "sample_test_doc.pdf"
+    generate_sample_pdf(pdf_path)
+    filename = pdf_path.name
+
+    # Ingestion
+    vector_store = qdrant.create_and_get_vector_storage(collection_name)
+    pipeline = IngestionPipeline(vector_store=vector_store)
+    pipeline.ingest(str(pdf_path), filename)
+
+    found = qdrant.collection_contains_file(
+        qdrant.client,
+        collection_name,
+        filename
+    )
+
+    assert found is True
+
+    # Clean the vector database
+    qdrant.delete_collection(collection_name)

genai/tests/integration/upload_test.py

@@ -0,0 +1,62 @@
+import io
+from unittest.mock import patch, MagicMock
+from fastapi.testclient import TestClient
+from main import app
+
+client = TestClient(app)
+
+
+@patch("routes.routes.qdrant.client.collection_exists", return_value=False)
+@patch("routes.routes.qdrant.create_and_get_vector_storage")
+@patch("routes.routes.IngestionPipeline")
+def test_upload_file_success(
+        mock_pipeline_class,
+        _mock_vector_store,
+        _mock_exists
+):
+    mock_pipeline = MagicMock()
+    mock_pipeline_class.return_value = mock_pipeline
+
+    file_content = b"%PDF-1.4 dummy content"
+    file = io.BytesIO(file_content)
+    file.name = "test.pdf"
+
+    response = client.post(
+        "/genai/upload",
+        files={"file": ("test.pdf", file, "application/pdf")}
+    )
+
+    assert response.status_code == 200
+    assert response.json() == {"message": "File processed successfully."}
+
+    mock_pipeline_class.assert_called_once()
+    mock_pipeline.ingest.assert_called_once()
+
+
+def test_upload_file_invalid_type():
+    file = io.BytesIO(b"just some text")
+    file.name = "notes.txt"
+
+    response = client.post(
+        "/genai/upload",
+        files={"file": ("notes.txt", file, "text/plain")}
+    )
+
+    assert response.status_code == 400
+    assert (response.json()["detail"] ==
+            "Invalid file type. Only PDF files are allowed.")
+
+
+@patch("routes.routes.qdrant.client.collection_exists", return_value=True)
+@patch("routes.routes.qdrant.collection_contains_file", return_value=True)
+def test_upload_file_already_exists(_mock_contains, _mock_exists):
+    file = io.BytesIO(b"%PDF-1.4")
+    file.name = "existing.pdf"
+
+    response = client.post(
+        "/genai/upload",
+        files={"file": ("existing.pdf", file, "application/pdf")}
+    )
+
+    assert response.status_code == 200
+    assert "already uploaded" in response.json()["message"]

genai/tests/unit/test_chat_model.py

@@ -0,0 +1,35 @@
+from rag.llm.chat_model import ChatModel
+from langchain_core.messages import HumanMessage, AIMessage
+from unittest.mock import patch
+
+
+def test_llm_type_property():
+    model = ChatModel()
+    assert model._llm_type == "recipai-custom-model"
+
+
+def test_get_system_prompt_contains_context():
+    model = ChatModel()
+    prompt = model.get_system_prompt()
+    assert isinstance(prompt, str)
+    assert "{context}" in prompt
+
+
+@patch("rag.llm.chat_model.generate_response")
+def test_generate_calls_openwebui_and_returns_response(mock_generate):
+    mock_generate.return_value = "This is a mock response"
+    model = ChatModel(model_name="mock-model")
+
+    messages = [
+        HumanMessage(content="What can I cook with potatoes?"),
+        AIMessage(content="You can make mashed potatoes."),
+        HumanMessage(content="Anything more creative?")
+    ]
+
+    result = model._generate(messages)
+    assert result.generations[0].message.content == "This is a mock response"
+    mock_generate.assert_called_once()
+    called_model_name, called_prompt = mock_generate.call_args[0]
+    assert called_model_name == "mock-model"
+    assert "User: What can I cook with potatoes?" in called_prompt
+    assert "Assistant: You can make mashed potatoes." in called_prompt

genai/tests/unit/test_ingestion_pipeline.py

@@ -0,0 +1,65 @@
+from unittest.mock import MagicMock, patch
+from uuid import UUID
+from rag.ingestion_pipeline import IngestionPipeline
+from langchain_core.documents import Document
+
+
+def test_load_document_returns_documents():
+    with patch("rag.ingestion_pipeline.PyPDFLoader") as mock_loader:
+        mock_loader.return_value.load.return_value = [
+            Document(page_content="Test")
+        ]
+        pipeline = IngestionPipeline(vector_store=MagicMock())
+        docs = pipeline.load_document("fake_path.pdf")
+        assert isinstance(docs, list)
+        assert isinstance(docs[0], Document)
+        assert docs[0].page_content == "Test"
+
+
+def test_chunk_documents_returns_chunks():
+    pipeline = IngestionPipeline(vector_store=MagicMock())
+    dummy_doc = Document(
+        page_content="This is a long text. " * 100,
+        metadata={}
+    )
+    chunks = pipeline.chunk_documents([dummy_doc], filename="sample.pdf")
+    assert isinstance(chunks, list)
+    assert all(isinstance(doc, Document) for doc in chunks)
+    assert all(doc.metadata["source"] == "sample.pdf" for doc in chunks)
+
+
+def test_store_documents_calls_add_documents_with_uuids():
+    mock_vector_store = MagicMock()
+    pipeline = IngestionPipeline(vector_store=mock_vector_store)
+    docs = [Document(page_content="Chunk", metadata={}) for _ in range(3)]
+    pipeline.store_documents(docs)
+    args, kwargs = mock_vector_store.add_documents.call_args
+    passed_docs = args[0]
+    passed_ids = kwargs["ids"]
+    assert len(passed_docs) == 3
+    assert len(passed_ids) == 3
+    assert all(UUID(uid) for uid in passed_ids)
+
+
+def test_ingest_calls_all_steps():
+    pipeline = IngestionPipeline(vector_store=MagicMock())
+
+    with patch.object(pipeline, "load_document") as mock_load, \
+         patch.object(pipeline, "chunk_documents") as mock_chunk, \
+         patch.object(pipeline, "store_documents") as mock_store, \
+         patch("rag.ingestion_pipeline.logger") as mock_logger, \
+         patch("rag.ingestion_pipeline.file_ingestion_duration.observe"), \
+         patch("rag.ingestion_pipeline.file_ingested_counter.inc"):
+
+        mock_load.return_value = [Document(page_content="Doc")]
+        mock_chunk.return_value = [Document(page_content="Chunk")]
+
+        pipeline.ingest("test.pdf", "testfile.pdf")
+
+        mock_load.assert_called_once_with("test.pdf")
+        mock_chunk.assert_called_once()
+        mock_store.assert_called_once()
+        mock_logger.info.assert_any_call(
+            "Documents are loaded for file %s",
+            "testfile.pdf"
+        )

genai/tests/unit/test_prompt.py

@@ -0,0 +1,50 @@
+from service.rag_service import process_raw_messages, prepare_prompt
+from langchain_core.messages import HumanMessage, AIMessage
+
+
+def test_process_raw_messages_creates_correct_types():
+    raw = [
+        {"role": "user", "content": "Hi"},
+        {"role": "assistant", "content": "Hello!"}
+    ]
+    messages = process_raw_messages(raw)
+    assert isinstance(messages[0], HumanMessage)
+    assert isinstance(messages[1], AIMessage)
+    assert messages[0].content == "Hi"
+    assert messages[1].content == "Hello!"
+
+
+def test_process_raw_messages_ignores_unknown_roles():
+    raw = [
+        {"role": "user", "content": "Hi"},
+        {"role": "system", "content": "Should be ignored"}
+    ]
+    messages = process_raw_messages(raw)
+    assert len(messages) == 1
+    assert isinstance(messages[0], HumanMessage)
+
+
+def test_prepare_prompt_structure():
+    system_prompt = "You are a helpful assistant. Context: {context}"
+    query = "What's a good recipe with eggs?"
+    docs = "Here are some egg-based recipes."
+    messages = [AIMessage(content="Hi there!")]
+
+    prompt = prepare_prompt(system_prompt, query, docs, messages)
+
+    assert hasattr(prompt, "to_messages")
+    final_messages = prompt.to_messages()
+    assert isinstance(final_messages[-1], HumanMessage)
+    assert "What's a good recipe with eggs?" in final_messages[-1].content
+
+
+def test_prepare_prompt_includes_docs_context():
+    system_prompt = "Use this: {context}"
+    query = "Tell me something"
+    docs = "Documented info"
+    messages = []
+
+    prompt = prepare_prompt(system_prompt, query, docs, messages)
+    rendered = prompt.to_string()
+    assert "Documented info" in rendered
+    assert "Tell me something" in rendered