From 22178e8a808ec3def1f6a90b992336d74c31d06c Mon Sep 17 00:00:00 2001
From: NeurArk <guillaume.rospape@neurark.com>
Date: Tue, 27 May 2025 17:22:57 +0200
Subject: [PATCH] test: add demo flow and integration tests

---
 TODO.md                   | 32 ++++++++--------
 tests/test_demo_flow.py   | 80 +++++++++++++++++++++++++++++++++++++++
 tests/test_integration.py | 17 +++++++++
 3 files changed, 113 insertions(+), 16 deletions(-)
 create mode 100644 tests/test_demo_flow.py
 create mode 100644 tests/test_integration.py

diff --git a/TODO.md b/TODO.md
index 38276f2..5077922 100644
--- a/TODO.md
+++ b/TODO.md
@@ -376,26 +376,26 @@ core/
 **Goal**: Ensure demo works flawlessly every time
 
 ### Tasks:
-- [ ] **Core Tests**
-  - [ ] Test chat engine with mock responses
-  - [ ] Test document upload flow
-  - [ ] Test RAG pipeline integration
-  - [ ] Ensure 75%+ coverage maintained
-
-- [ ] **Demo Scenario Tests**
-  - [ ] Test with sample PDFs
-  - [ ] Test common questions
-  - [ ] Test error cases (bad file, no context)
-  - [ ] Full demo run-through test
+- [x] **Core Tests**
+  - [x] Test chat engine with mock responses
+  - [x] Test document upload flow
+  - [x] Test RAG pipeline integration
+  - [x] Ensure 75%+ coverage maintained
+
+- [x] **Demo Scenario Tests**
+  - [x] Test with sample PDFs
+  - [x] Test common questions
+  - [x] Test error cases (bad file, no context)
+  - [x] Full demo run-through test
 
 ### Acceptance Criteria:
-- [ ] All tests pass reliably
-- [ ] Demo scenarios fully covered
-- [ ] No flaky tests
+- [x] All tests pass reliably
+- [x] Demo scenarios fully covered
+- [x] No flaky tests
 
 ### Definition of Done:
-- [ ] Tests ensure demo reliability
-- [ ] CI/CD passing consistently
+- [x] Tests ensure demo reliability
+- [x] CI/CD passing consistently
 
 ---
 
diff --git a/tests/test_demo_flow.py b/tests/test_demo_flow.py
new file mode 100644
index 0000000..5b190a7
--- /dev/null
+++ b/tests/test_demo_flow.py
@@ -0,0 +1,80 @@
+import pytest
+from unittest.mock import Mock, patch
+from pathlib import Path
+import fitz
+
+from core.document_processor import DocumentProcessor
+from core.embedder import EmbeddingService
+from core.rag_pipeline import RAGPipeline
+
+
+@pytest.fixture()
+def sample_pdf(tmp_path: Path) -> str:
+    pdf_path = tmp_path / "test.pdf"
+    doc = fitz.open()
+    page = doc.new_page()
+    page.insert_text((72, 72), "Hello World")
+    doc.save(str(pdf_path))
+    return str(pdf_path)
+
+
+@pytest.fixture()
+def mock_openai() -> Mock:
+    with patch("openai.OpenAI") as mock_cls:
+        client = Mock()
+        mock_cls.return_value = client
+
+        embed_resp = Mock()
+        embed_resp.data = [Mock(embedding=[0.1] * 3072)]
+        client.embeddings.create.return_value = embed_resp
+
+        chat_resp = Mock()
+        chat_resp.choices = [Mock(message=Mock(content="Test answer"))]
+        client.chat.completions.create.return_value = chat_resp
+        yield client
+
+
+def test_complete_demo_flow(sample_pdf: str, mock_openai: Mock, tmp_path, monkeypatch) -> None:
+    monkeypatch.setenv("CHROMA_PERSIST_DIR", str(tmp_path))
+    from config.settings import get_settings
+    get_settings.cache_clear()
+    import importlib
+    import core.vector_store as vs
+    importlib.reload(vs)
+
+    processor = DocumentProcessor()
+    embedder = EmbeddingService()
+    vector_store = vs.VectorStore()
+    rag = RAGPipeline()
+
+    doc, chunks = processor.process_document(sample_pdf)
+    assert doc is not None
+    assert len(chunks) > 0
+
+    embedded = embedder.embed_document(doc, chunks)
+    assert all(c.embedding is not None for c in embedded)
+
+    vector_store.store_document(doc, embedded)
+
+    answer, sources = rag.query("What is this document about?")
+    assert len(answer) > 0
+    assert len(sources) >= 0
+
+
+def test_no_document_handling(mock_openai: Mock, tmp_path, monkeypatch) -> None:
+    monkeypatch.setenv("CHROMA_PERSIST_DIR", str(tmp_path))
+    from config.settings import get_settings
+    get_settings.cache_clear()
+    import importlib
+    import core.vector_store as vs
+    importlib.reload(vs)
+    rag = RAGPipeline()
+    answer, sources = rag.query("Tell me about the contract")
+    assert "couldn't find" in answer.lower()
+    assert sources == []
+
+
+def test_error_handling(mock_openai: Mock) -> None:
+    processor = DocumentProcessor()
+    with pytest.raises(Exception):
+        processor.process_document("does_not_exist.pdf")
diff --git a/tests/test_integration.py b/tests/test_integration.py
new file mode 100644
index 0000000..7adfc59
--- /dev/null
+++ b/tests/test_integration.py
@@ -0,0 +1,17 @@
+import pytest
+from unittest.mock import patch, Mock
+
+import app
+
+
+def test_gradio_functions():
+    with patch.object(app.vector_store, "clear", return_value=None) as clear_mock:
+        result = app.clear_all_documents()
+        assert "cleared" in result.lower()
+        clear_mock.assert_called_once()
+
+    with patch.object(app.rag_pipeline, "query", return_value=("Hello", [])) as qmock:
+        response = app.chat_response("Hi", [])
+        assert len(response) > 0
+        assert "error" not in response.lower()
+        qmock.assert_called_once()