catalog generation unit tests fixed

Author: tolgaerdonmez

backend/llm_eval/qa_catalog/generator/implementation/ragas/generator.py

@@ -303,6 +303,7 @@ def load_exiting_knowledge_graph(
             return None
 
         logger.info(f"Knowledge graph found at {self.config.knowledge_graph_location}")
+        knowledge_graph = None
         try:
             knowledge_graph = KnowledgeGraph.load(self.config.knowledge_graph_location)
         except Exception as e:

backend/tests/unit/backend/qa_catalog/logic/test_generation.py

@@ -124,12 +124,8 @@ def catalog_generation_data(
 def session_local() -> MockSessionLocal:
     @contextmanager
     def _func():  # noqa: ANN202
-        with patch(
-            "llm_eval.qa_catalog.logic.generation.AsyncSessionLocal",
-        ) as session_local:
-            mock_session = AsyncMock()
-            session_local.begin.return_value.__aenter__.return_value = mock_session
-            yield mock_session
+        mock_session = AsyncMock()
+        yield mock_session
 
     return _func
 
@@ -153,7 +149,7 @@ async def test_generate_catalog_task_happy_path(
         mock_find_qa_catalog.return_value = qa_catalog
         mock_find_data_source_config.return_value = temp_data_source_config
 
-        await generate_catalog(qa_catalog.id, catalog_generation_data)
+        await generate_catalog(mock_session, qa_catalog.id, catalog_generation_data)
 
         mock_find_qa_catalog.assert_called_once_with(mock_session, qa_catalog.id)
         mock_find_data_source_config.assert_called_once_with(
@@ -194,7 +190,7 @@ async def test_generate_catalog_task_temp_data_source_config_not_found(
         mock_find_qa_catalog.return_value = qa_catalog
         mock_find_data_source_config.return_value = None
 
-        await generate_catalog(qa_catalog.id, catalog_generation_data)
+        await generate_catalog(mock_session, qa_catalog.id, catalog_generation_data)
 
         mock_find_qa_catalog.assert_called_once_with(mock_session, qa_catalog.id)
         mock_find_data_source_config.assert_called_once_with(

backend/tests/unit/qa_catalog/generator/implementation/ragas/test_generator.py

@@ -2,11 +2,9 @@
 from pathlib import Path
 from random import shuffle
 from typing import Callable, ContextManager, Generator
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 import pytest
-from anyio import CapacityLimiter
-from anyio.streams.memory import MemoryObjectSendStream
 from langchain_core.documents import Document
 from ragas.testset.graph import Node, NodeType
 
@@ -234,31 +232,27 @@ async def test_ragas_generator_has_same_nodes(
 
 
 @pytest.mark.asyncio
-@patch("llm_eval.qa_catalog.generator.implementation.ragas.generator.create_backup")
 @patch("llm_eval.qa_catalog.generator.implementation.ragas.generator.KnowledgeGraph")
 async def test_ragas_generator_create_knowledge_graph_new_successfully(
     mock_knowledge_graph: MagicMock,
-    mock_create_backup: MagicMock,
     ragas_generator_factory: RagasGeneratorFactory,
     documents: list[Document],
     nodes: list[Node],
 ) -> None:
     with ragas_generator_factory() as generator:
         generator._load_and_process_documents = MagicMock(return_value=documents)
         generator._create_knowledge_graph_nodes = MagicMock(return_value=nodes)
-        generator.load_exiting_knowledge_graph = MagicMock(return_value=None)
+        generator.split_documents = MagicMock(return_value=documents)
         generator.apply_knowledge_graph_transformations = MagicMock()
 
         created_kg = mock_knowledge_graph(nodes=nodes)
         mock_knowledge_graph.return_value = created_kg
 
         kg = generator.create_knowledge_graph()
 
+        generator._load_and_process_documents.assert_called_once()
+        generator.split_documents.assert_called_once_with(documents)
         generator._create_knowledge_graph_nodes.assert_called_once_with(documents)
-        generator.load_exiting_knowledge_graph.assert_called_once()
-        mock_create_backup.assert_called_once_with(
-            generator.config.knowledge_graph_location
-        )
         mock_knowledge_graph.assert_called_with(nodes=nodes)
         generator.apply_knowledge_graph_transformations.assert_called_once_with(
             created_kg
@@ -277,16 +271,16 @@ async def test_ragas_generator_create_knowledge_graph_use_existing_graph_when_no
     with ragas_generator_factory() as generator:
         generator._load_and_process_documents = MagicMock(return_value=documents)
         generator._create_knowledge_graph_nodes = MagicMock(return_value=nodes)
+        generator.split_documents = MagicMock(return_value=documents)
         generator.config.knowledge_graph_location = None
         generator.apply_knowledge_graph_transformations = MagicMock(return_value=None)
-        existing_graph = MagicMock(name="existing_graph", nodes=nodes)
-        generator.load_exiting_knowledge_graph = MagicMock(return_value=existing_graph)
-        mock_knowledge_graph.return_value = MagicMock(name="generated_graph")
+        generated_graph = MagicMock(name="generated_graph")
+        mock_knowledge_graph.return_value = generated_graph
 
         kg: MagicMock = generator.create_knowledge_graph()  # type: ignore
 
         assert kg is not None
-        assert kg == existing_graph
+        assert kg == generated_graph
 
 
 @pytest.mark.asyncio
@@ -310,14 +304,24 @@ async def test_ragas_generator_load_knowledge_graph(
     with ragas_generator_factory() as generator:
         generator.config.knowledge_graph_location = MagicMock()
         generator.config.knowledge_graph_location.exists = MagicMock(return_value=True)
-        mock_knowledge_graph.load.return_value = MagicMock(name="existing_graph")
+        loaded_graph = MagicMock(name="existing_graph")
+        mock_knowledge_graph.load.return_value = loaded_graph
+
+        # Create proper mock documents with required attributes
+        mock_doc = MagicMock()
+        mock_doc.page_content = "test content"
+        mock_doc.metadata = {"source": "test.txt"}
+        docs = [mock_doc]
 
-        kg = generator.load_exiting_knowledge_graph()
+        # Mock _has_same_nodes to return True so the loaded graph is returned
+        generator._has_same_nodes = MagicMock(return_value=True)
+
+        kg = generator.load_exiting_knowledge_graph(docs)
 
         mock_knowledge_graph.load.assert_called_once_with(
             generator.config.knowledge_graph_location
         )
-        assert kg is not None
+        assert kg is loaded_graph
 
 
 @pytest.mark.asyncio
@@ -330,7 +334,8 @@ async def test_ragas_generator_load_knowledge_graph_fails_on_non_existent_file(
         generator.config.knowledge_graph_location = MagicMock()
         generator.config.knowledge_graph_location.exists = MagicMock(return_value=False)
 
-        kg = generator.load_exiting_knowledge_graph()
+        docs = [MagicMock()]  # Mock documents
+        kg = generator.load_exiting_knowledge_graph(docs)
 
         assert kg is None
 
@@ -342,7 +347,8 @@ async def test_ragas_generator_load_knowledge_graph_fails_on_none_location(
     with ragas_generator_factory() as generator:
         generator.config.knowledge_graph_location = None
 
-        kg = generator.load_exiting_knowledge_graph()
+        docs = [MagicMock()]  # Mock documents
+        kg = generator.load_exiting_knowledge_graph(docs)
 
         assert kg is None
 
@@ -362,7 +368,13 @@ async def test_ragas_generator_load_knowledge_graph_fails_on_load(
         generator.config.knowledge_graph_location.exists = MagicMock(return_value=True)
         mock_knowledge_graph.load.side_effect = RuntimeError("Load failed")
 
-        kg = generator.load_exiting_knowledge_graph()
+        # Create proper mock documents with required attributes
+        mock_doc = MagicMock()
+        mock_doc.page_content = "test content"
+        mock_doc.metadata = {"source": "test.txt"}
+        docs = [mock_doc]
+
+        kg = generator.load_exiting_knowledge_graph(docs)
 
         mock_knowledge_graph.load.assert_called_once_with(
             generator.config.knowledge_graph_location
@@ -382,117 +394,72 @@ async def test_ragas_generator_generate_testset(
 
         mock_testset_generator = MagicMock()
         mock_testset_generator.generate.return_value = ["generated query"]
-        testset = generator.generate_testset(mock_testset_generator, 1)
-
-        assert testset[0] == "generated query"  # type: ignore
-
-
-@pytest.mark.asyncio
-@patch(
-    "llm_eval.qa_catalog.generator.implementation.ragas.generator.ragas_sample_to_synthetic_qa_pair",
-)
-async def test_ragas_generator_generate_single_sample(
-    mock_from_ragas: MagicMock,
-    ragas_generator_factory: RagasGeneratorFactory,
-) -> None:
-    with ragas_generator_factory() as generator:
-        generator.generate_testset = MagicMock(return_value=MagicMock(samples=[1]))
-        mock_from_ragas.return_value = "sample"
-
-        send_sample = MagicMock()
-        send_sample.send = AsyncMock()
-        limiter = MagicMock()
-
-        await generator._generate_samples(MagicMock(), send_sample, limiter)
 
-        send_sample.send.assert_called_once_with("sample")
+        # Provide a non-empty query distribution
+        mock_synthesizer = MagicMock()
+        query_distribution = [(mock_synthesizer, 1.0)]
 
-
-@pytest.mark.asyncio
-async def test_ragas_generator_generate_single_sample_wont_send_when_fails(
-    ragas_generator_factory: RagasGeneratorFactory,
-) -> None:
-    with ragas_generator_factory() as generator:
-        generator.generate_testset = MagicMock(
-            side_effect=RuntimeError("sample generation failed")
+        testset = generator.generate_testset(
+            mock_testset_generator, 1, query_distribution
         )
 
-        send_sample = MagicMock()
-        send_sample.send = AsyncMock()
-        limiter = MagicMock()
-
-        await generator._generate_samples(MagicMock(), send_sample, limiter)
-
-        send_sample.send.assert_not_called()
-
-
-@pytest.mark.asyncio
-async def test_ragas_generator_generate_single_sample_does_nothing_on_empty_generated_testset(  # noqa: E501
-    ragas_generator_factory: RagasGeneratorFactory,
-) -> None:
-    with ragas_generator_factory() as generator:
-        generator.generate_testset = MagicMock(return_value=MagicMock(samples=[]))
-
-        send_sample = MagicMock()
-        send_sample.send = AsyncMock()
-        limiter = MagicMock()
-
-        await generator._generate_samples(MagicMock(), send_sample, limiter)
-
-        send_sample.send.assert_not_called()
+        assert testset[0] == "generated query"  # type: ignore
 
 
 @pytest.mark.asyncio
 @patch(
-    "llm_eval.qa_catalog.generator.implementation.ragas.generator.TestsetGenerator",
-    new_callable=AsyncMock,
-)
-@patch(
-    "llm_eval.qa_catalog.generator.implementation.ragas.generator.RagasQACatalogGenerator.load_chat_model",
-    new_callable=AsyncMock,
-)
-@patch(
-    "llm_eval.qa_catalog.generator.implementation.ragas.generator.copy",
-    new_callable=AsyncMock,
+    "llm_eval.qa_catalog.generator.implementation.ragas.generator.ragas_sample_to_synthetic_qa_pair"
 )
+@patch("ragas.testset.persona.generate_personas_from_kg")
 async def test_ragas_generator_a_create_synthetic_qa_successfull(
-    mock_copy: AsyncMock,
-    mock_load_chat_model: AsyncMock,
-    mock_testset_generator: AsyncMock,
-    config: RagasQACatalogGeneratorConfig,
-    ragas_model_config: RagasQACatalogGeneratorModelConfig,
-    data_source_config: QACatalogGeneratorDataSourceConfig,
+    mock_generate_personas: MagicMock,
+    mock_ragas_sample_to_qa_pair: MagicMock,
+    ragas_generator_factory: RagasGeneratorFactory,
 ) -> None:
-    generator = RagasQACatalogGenerator(config, data_source_config, ragas_model_config)
-    generator.create_knowledge_graph = AsyncMock()
-
-    sample = SyntheticQAPair(
-        id="1",
-        question="question",
-        expected_output="expected_output",
-        contexts=[],
-        meta_data={},
-    )
-
-    async def mock_generate(
-        _,  # noqa: ANN001
-        send_sample: MemoryObjectSendStream[SyntheticQAPair | None],
-        limiter: CapacityLimiter,
-    ) -> None:
-        async with limiter:
-            async with send_sample:
-                await send_sample.send(sample)
+    with ragas_generator_factory() as generator:
+        # Mock knowledge graph
+        mock_kg = MagicMock()
+        mock_kg.nodes = []
+        generator.create_knowledge_graph = MagicMock(return_value=mock_kg)
+
+        # Mock personas
+        mock_persona = MagicMock()
+        mock_generate_personas.return_value = [mock_persona]
+
+        # Mock query distribution
+        mock_synthesizer = MagicMock()
+        generator.create_query_distribution = MagicMock(
+            return_value=[(mock_synthesizer, 1.0)]
+        )
 
-    generator._generate_samples = AsyncMock(side_effect=mock_generate)
+        # Mock testset generation
+        mock_testset = MagicMock()
+        mock_sample = MagicMock()
+        mock_testset.samples = [mock_sample]
+        generator.generate_testset = MagicMock(return_value=mock_testset)
+
+        # Mock conversion to SyntheticQAPair
+        synthetic_qa_pair = SyntheticQAPair(
+            id="1",
+            question="test question",
+            expected_output="test answer",
+            contexts=[],
+            meta_data={},
+        )
+        mock_ragas_sample_to_qa_pair.return_value = synthetic_qa_pair
 
-    samples = []
+        # Collect samples
+        collected_samples = []
 
-    async def process_samples_fn(samples_batch: list[SyntheticQAPair]) -> None:
-        samples.extend(samples_batch)
+        async def collect_samples_fn(samples: list[SyntheticQAPair]) -> None:
+            collected_samples.extend(samples)
 
-    await generator.a_create_synthetic_qa(process_samples_fn)
+        # Test the method
+        await generator.a_create_synthetic_qa(collect_samples_fn)
 
-    assert len(samples) == generator.config.sample_count
+        # Verify results
+        assert len(collected_samples) == 1
+        assert collected_samples[0] == synthetic_qa_pair
 
 
 @pytest.mark.asyncio