refactor: enhance docstrings for mock classes and unit tests in retriever module

Author: a-klos

libs/rag-core-api/tests/composite_retriever_test.py

@@ -1,4 +1,4 @@
-"""Unit tests for internal helper methods of ``CompositeRetriever``.
+"""Test internal helper methods of ``CompositeRetriever``.
 
 The goal of these tests is to verify the transformation semantics of:
  - _use_summaries
@@ -40,6 +40,10 @@ def _mk_doc(
 
 @pytest.mark.asyncio
 async def test_use_summaries_expands_and_removes_summary():
+    """Expand a summary into its related documents.
+
+    Verify that summary documents are removed and replaced by their related underlying documents.
+    """
     # Summary references an underlying doc not in initial results.
     underlying = _mk_doc("doc1", score=0.9)
     summary = _mk_doc("sum1", doc_type=ContentType.SUMMARY, related=["doc1"])  # type: ignore[arg-type]
@@ -57,6 +61,10 @@ async def test_use_summaries_expands_and_removes_summary():
 
 
 def test_use_summaries_only_summary_no_related():
+    """Drop a summary document that has no related documents.
+
+    Verify that the returned result is empty when no related ids are present.
+    """
     summary = _mk_doc("sum1", doc_type=ContentType.SUMMARY, related=[])  # type: ignore[arg-type]
     retriever = MockRetrieverQuark([summary])
     cr = CompositeRetriever(retrievers=[retriever], reranker=None, reranker_enabled=False)
@@ -66,6 +74,10 @@ def test_use_summaries_only_summary_no_related():
 
 
 def test_remove_duplicates_preserves_first_occurrence():
+    """Preserve the first occurrence when duplicate ids are present.
+
+    Verify that duplicate documents are removed while maintaining the original order.
+    """
     d1a = _mk_doc("a")
     d1b = _mk_doc("a")  # duplicate id
     d2 = _mk_doc("b")
@@ -76,6 +88,10 @@ def test_remove_duplicates_preserves_first_occurrence():
 
 
 def test_early_pruning_sorts_by_score_when_all_have_score():
+    """Sort by score and keep only the top-k documents.
+
+    Verify that documents are sorted descending by score when all documents include scores.
+    """
     docs = [_mk_doc("a", score=0.7), _mk_doc("b", score=0.9), _mk_doc("c", score=0.8)]
     retriever = MockRetrieverQuark(docs)
     cr = CompositeRetriever(
@@ -87,6 +103,10 @@ def test_early_pruning_sorts_by_score_when_all_have_score():
 
 
 def test_early_pruning_preserves_order_without_scores():
+    """Preserve input order when pruning without score metadata.
+
+    Verify that pruning keeps the original order when scores are absent.
+    """
     docs = [_mk_doc("a"), _mk_doc("b"), _mk_doc("c")]  # no scores
     retriever = MockRetrieverQuark(docs)
     cr = CompositeRetriever(
@@ -98,6 +118,10 @@ def test_early_pruning_preserves_order_without_scores():
 
 @pytest.mark.asyncio
 async def test_arerank_pruning_invokes_reranker_when_needed():
+    """Invoke the reranker when more than k documents are retrieved.
+
+    Verify that the reranker is called and that the returned list is trimmed to ``reranker_k_documents``.
+    """
     docs = [_mk_doc("a", score=0.5), _mk_doc("b", score=0.7), _mk_doc("c", score=0.9)]
     retriever = MockRetrieverQuark(docs)
     reranker = MockReranker()
@@ -116,6 +140,10 @@ async def test_arerank_pruning_invokes_reranker_when_needed():
 
 @pytest.mark.asyncio
 async def test_arerank_pruning_skips_when_not_needed():
+    """Skip reranking when the retrieved docs are already within k.
+
+    Verify that the reranker is not invoked when no pruning is required.
+    """
     docs = [_mk_doc("a", score=0.5), _mk_doc("b", score=0.7)]  # already <= k
     retriever = MockRetrieverQuark(docs)
     reranker = MockReranker()

libs/rag-core-api/tests/mocks/mock_reranker.py

@@ -1,15 +1,32 @@
-"""Mock reranker used by CompositeRetriever unit tests."""
-
-from langchain_core.documents import Document
+"""Provide a mock reranker for CompositeRetriever unit tests."""
 
 __all__ = ["MockReranker"]
 
 
 class MockReranker:
+    """Provide a simple reranker test double.
+
+    The mock records whether it was invoked and returns a deterministic top-2 subset.
+    """
+
     def __init__(self):
         self.invoked = False
 
-    async def ainvoke(self, payload, config=None):  # noqa: D401
+    async def ainvoke(self, payload, config=None):
+        """Return a reranked subset of the provided documents.
+
+        Parameters
+        ----------
+        payload : tuple
+            A ``(documents, query)`` tuple.
+        config : Any, optional
+            Optional runtime config passed through by the caller.
+
+        Returns
+        -------
+        list
+            The top two documents sorted by score when available.
+        """
         self.invoked = True
         documents, _query = payload
         # Emulate reranker selecting top 2 with highest 'score' if present; else first 2 reversed

libs/rag-core-api/tests/mocks/mock_retriever_quark.py

@@ -1,6 +1,5 @@
-"""Mock retriever quark for CompositeRetriever unit tests."""
+"""Provide a mock retriever quark for CompositeRetriever unit tests."""
 
-from typing import List
 from langchain_core.documents import Document
 
 from .mock_vector_db import MockVectorDB
@@ -9,18 +8,31 @@
 
 
 class MockRetrieverQuark:
-    """Minimal stand-in for a RetrieverQuark.
+    """Provide a minimal stand-in for a RetrieverQuark.
 
     Exposes an ``ainvoke`` returning pre-seeded documents and a ``_vector_database`` attribute
     referenced by summary expansion logic.
     """
 
-    def __init__(self, documents: List[Document], vector_database: MockVectorDB | None = None):
+    def __init__(self, documents: list[Document], vector_database: MockVectorDB | None = None):
         self._documents = documents
         self._vector_database = vector_database or MockVectorDB()
 
     def verify_readiness(self):  # pragma: no cover - trivial
-        return None
-
-    async def ainvoke(self, *_args, **_kwargs):  # noqa: D401 - simple passthrough
+        """Verify that the retriever is ready.
+
+        Returns
+        -------
+        None
+            Always returns ``None``.
+        """
+
+    async def ainvoke(self, *_args, **_kwargs):
+        """Return the pre-seeded documents.
+
+        Returns
+        -------
+        list[Document]
+            The documents passed to the constructor.
+        """
         return self._documents

libs/rag-core-api/tests/mocks/mock_vector_db.py

@@ -1,23 +1,43 @@
-"""Mock implementation of a minimal vector database interface for tests.
+"""Provide a minimal vector database interface for tests.
 
 Provides only the methods required by the CompositeRetriever unit tests:
 - get_documents_by_ids: Used during summary expansion
 - asearch: (async) provided as a defensive stub
 """
 
-from typing import Dict, List
 from langchain_core.documents import Document
 
 __all__ = ["MockVectorDB"]
 
 
 class MockVectorDB:
-    def __init__(self, docs_by_id: Dict[str, Document] | None = None):
+    """Provide a minimal in-memory vector database test double."""
+
+    def __init__(self, docs_by_id: dict[str, Document] | None = None):
         self.collection_available = True
         self._docs_by_id = docs_by_id or {}
 
-    def get_documents_by_ids(self, ids: List[str]) -> List[Document]:  # pragma: no cover - simple mapping
+    def get_documents_by_ids(self, ids: list[str]) -> list[Document]:  # pragma: no cover - simple mapping
+        """Return documents for the provided ids.
+
+        Parameters
+        ----------
+        ids : list[str]
+            Document ids to look up.
+
+        Returns
+        -------
+        list[Document]
+            Documents that exist in the in-memory mapping.
+        """
         return [self._docs_by_id[i] for i in ids if i in self._docs_by_id]
 
     async def asearch(self, *_, **__):  # pragma: no cover - defensive stub
+        """Return an empty result for async search.
+
+        Returns
+        -------
+        list
+            Always returns an empty list.
+        """
         return []