test: integration test robustness — stubs, retries, async, diagnostics

Makefile

@@ -30,7 +30,7 @@ test-integration:
 	$(PYTEST_INTEGRATION) --maxfail=1
 
 test-integration-verbose:
-	uv run pytest tests/integration -v -rA --durations=10
+	$(PYTEST_INTEGRATION) -s -rA --durations=10 --tb=long
 
 test-integration-openai:
 	$(PYTEST_INTEGRATION) -m openai

deepfabric/generator.py

@@ -1461,10 +1461,10 @@ async def _run_generation_loop_async(  # noqa: PLR0912, PLR0915
         topic_model_type: str | None = None,
     ) -> AsyncGenerator[dict | HFDataset, None]:
         """Run the main generation loop yielding progress events."""
-        # Verify topic paths cover all expected samples
+        # Verify topic paths cover all expected samples (only when a topic model is used)
         expected_prompts = num_steps * batch_size
         actual_paths = len(topic_paths) if topic_paths else 0
-        if actual_paths < expected_prompts:
+        if topic_paths and actual_paths < expected_prompts:
             logger.warning(
                 "Topic paths (%d) < expected samples (%d). Steps beyond path %d will produce 0 samples.",
                 actual_paths,

pyproject.toml

@@ -34,6 +34,8 @@ dev = [
     "pytest>=7.0.0",
     "pytest-cov>=4.0.0",
     "pytest-mock>=3.10.0",
+    "pytest-rerunfailures>=14.0",
+    "pytest-asyncio>=0.23.0",
     "pytest-httpx>=0.30.0",
     "requests-mock>=1.11.0",
     "ruff>=0.1.0",
@@ -74,12 +76,14 @@ filterwarnings = [
     "ignore::ImportWarning",
     "ignore::pydantic.PydanticDeprecatedSince20",
 ]
+asyncio_mode = "auto"
 markers = [
     "asyncio: mark tests that require an asyncio event loop",
     "openai: mark tests requiring OpenAI API access",
     "gemini: mark tests requiring Gemini API access",
     "huggingface: mark tests requiring HuggingFace Hub access",
     "spin: mark tests for Spin service (mocked HTTP)",
+    "flaky: mark tests that may fail due to external LLM service variability",
 ]
 
 [tool.ruff]

tests/integration/conftest.py

@@ -1,9 +1,18 @@
 """Shared fixtures and markers for integration tests."""
 
+import logging
 import os
 
 import pytest
 
+# Suppress httpx async client cleanup noise.
+# LLMClient doesn't expose aclose(), so httpx connections linger after
+# pytest-asyncio tears down the event loop, producing harmless
+# "Event loop is closed" ERROR logs.
+logging.getLogger("asyncio").addFilter(
+    lambda record: "Event loop is closed" not in record.getMessage()
+)
+
 # Skip markers for conditional test execution based on API key availability
 requires_openai = pytest.mark.skipif(
     not os.getenv("OPENAI_API_KEY"),
@@ -40,3 +49,63 @@ def gemini_config():
         "model_name": os.getenv("GEMINI_TEST_MODEL", "gemini-2.0-flash"),
         "temperature": 0.2,
     }
+
+
+def assert_generation_result(result, generator, min_samples=1, context=""):
+    """Assert generation produced enough samples, with diagnostics on failure.
+
+    Args:
+        result: HuggingFace Dataset returned by create_data.
+        generator: DataSetGenerator instance for reading failure diagnostics.
+        min_samples: Minimum expected sample count.
+        context: Test context for the failure message.
+    """
+    if len(result) >= min_samples:
+        return
+
+    lines = [
+        f"Generation produced {len(result)} samples, expected >= {min_samples}.",
+    ]
+    if context:
+        lines.append(f"Context: {context}")
+
+    lines.append(f"Total failures: {len(generator.failed_samples)}")
+
+    failure_counts = {k: len(v) for k, v in generator.failure_analysis.items() if v}
+    if failure_counts:
+        lines.append("Failure breakdown:")
+        for ftype, count in sorted(failure_counts.items()):
+            lines.append(f"  {ftype}: {count}")
+
+    for i, f in enumerate(generator.failed_samples[:5], 1):
+        error_msg = f.get("error", str(f)) if isinstance(f, dict) else str(f)
+        lines.append(f"  [{i}] {error_msg[:300]}")
+
+    pytest.fail("\n".join(lines))
+
+
+def assert_topic_build_result(paths, model, min_paths=1, context=""):
+    """Assert topic model build produced enough paths, with diagnostics on failure.
+
+    Args:
+        paths: List of paths from get_all_paths().
+        model: Tree or Graph instance for reading failed_generations.
+        min_paths: Minimum expected path count.
+        context: Test context for the failure message.
+    """
+    if len(paths) >= min_paths:
+        return
+
+    lines = [
+        f"Topic build produced {len(paths)} paths, expected >= {min_paths}.",
+    ]
+    if context:
+        lines.append(f"Context: {context}")
+
+    failed_gens = getattr(model, "failed_generations", [])
+    if failed_gens:
+        lines.append(f"Failed generations ({len(failed_gens)}):")
+        for i, fg in enumerate(failed_gens[:5], 1):
+            lines.append(f"  [{i}] {fg}")
+
+    pytest.fail("\n".join(lines))

tests/integration/test_generator_integration.py

@@ -6,7 +6,12 @@
 
 from deepfabric import DataSetGenerator, Graph
 
-from .conftest import requires_gemini, requires_openai
+from .conftest import (
+    assert_generation_result,
+    assert_topic_build_result,
+    requires_gemini,
+    requires_openai,
+)
 
 
 @pytest.fixture
@@ -48,27 +53,32 @@ async def build_graph():
             pass
 
     asyncio.run(build_graph())
+    assert_topic_build_result(
+        graph.get_all_paths(), graph, min_paths=1, context="small_topic_graph fixture"
+    )
     return graph
 
 
+@requires_openai
 class TestDataSetGeneratorOpenAI:
     """Integration tests for DataSetGenerator with OpenAI provider."""
 
-    @requires_openai
     @pytest.mark.openai
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
     def test_basic_generation(self, openai_generator):
         """Test basic dataset generation without topic model."""
         result = openai_generator.create_data(
             num_steps=1,
             batch_size=2,
         )
 
-        # Result should be a HuggingFace Dataset
         assert result is not None
-        assert len(result) >= 1
+        assert_generation_result(
+            result, openai_generator, context="basic generation (no topic model)"
+        )
 
-    @requires_openai
     @pytest.mark.openai
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
     def test_generation_with_topic_model(self, openai_generator, small_topic_graph):
         """Test dataset generation with a topic graph."""
         result = openai_generator.create_data(
@@ -78,26 +88,22 @@ def test_generation_with_topic_model(self, openai_generator, small_topic_graph):
         )
 
         assert result is not None
-        assert len(result) >= 1
+        assert_generation_result(result, openai_generator, context="generation with topic model")
 
-    @requires_openai
     @pytest.mark.openai
-    def test_async_generation(self, openai_generator):
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
+    async def test_async_generation(self, openai_generator):
         """Test async dataset generation."""
-
-        async def run_async():
-            return await openai_generator.create_data_async(
-                num_steps=1,
-                batch_size=2,
-            )
-
-        result = asyncio.run(run_async())
+        result = await openai_generator.create_data_async(
+            num_steps=1,
+            batch_size=2,
+        )
 
         assert result is not None
-        assert len(result) >= 1
+        assert_generation_result(result, openai_generator, context="async generation")
 
-    @requires_openai
     @pytest.mark.openai
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
     def test_generation_with_cot(self, openai_config):
         """Test generation with cot conversation type."""
         generator = DataSetGenerator(
@@ -115,40 +121,33 @@ def test_generation_with_cot(self, openai_config):
         )
 
         assert result is not None
-        assert len(result) >= 1
+        assert_generation_result(result, generator, context="cot generation")
 
 
+@requires_gemini
 class TestDataSetGeneratorGemini:
     """Integration tests for DataSetGenerator with Gemini provider."""
 
-    @requires_gemini
     @pytest.mark.gemini
-    def test_basic_generation(self, gemini_generator):
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
+    async def test_basic_generation(self, gemini_generator):
         """Test basic dataset generation with Gemini."""
-
-        async def run_async():
-            return await gemini_generator.create_data_async(
-                num_steps=1,
-                batch_size=2,
-            )
-
-        result = asyncio.run(run_async())
+        result = await gemini_generator.create_data_async(
+            num_steps=1,
+            batch_size=2,
+        )
 
         assert result is not None
-        assert len(result) >= 1
+        assert_generation_result(result, gemini_generator, context="Gemini basic generation")
 
-    @requires_gemini
     @pytest.mark.gemini
-    def test_generation_saves_to_file(self, tmp_path, gemini_generator):
+    @pytest.mark.flaky(reruns=2, reruns_delay=5)
+    async def test_generation_saves_to_file(self, tmp_path, gemini_generator):
         """Test that generated data can be saved."""
-
-        async def run_async():
-            return await gemini_generator.create_data_async(
-                num_steps=1,
-                batch_size=2,
-            )
-
-        asyncio.run(run_async())
+        await gemini_generator.create_data_async(
+            num_steps=1,
+            batch_size=2,
+        )
 
         # Save dataset
         out_path = tmp_path / "dataset.jsonl"