Add tests demonstrating decorator generator handling

Author: mmacpherson

pyproject.toml

@@ -52,6 +52,8 @@ dev = [
   "fastapi>=0.116.1",
   "litestar>=2.17",
   "pre-commit>=4.2",
+  "pytest>=8.4.2",
+  "pytest-asyncio>=1.2.0",
   "python-fasthtml>=0.12.25; python_full_version>='3.10'",
   "quart>=0.20",
   "sanic>=25.3",

tests/conftest.py

@@ -0,0 +1,10 @@
+"""Test configuration."""
+
+import sys
+from pathlib import Path
+
+# Ensure the local src/ package is importable without installing
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))

tests/test_datastar_response.py

@@ -0,0 +1,128 @@
+"""Tests for datastar_response decorator with various function types."""
+
+import importlib.util
+
+import pytest
+
+from datastar_py.sse import ServerSentEventGenerator as SSE
+
+
+def _has(module: str) -> bool:
+    return importlib.util.find_spec(module) is not None
+
+
+pytest.importorskip("starlette")
+from datastar_py.starlette import DatastarResponse, datastar_response  # noqa: E402
+
+
+@pytest.mark.asyncio
+async def test_async_generator_support():
+    """Async generators should yield multiple events."""
+
+    @datastar_response
+    async def multi_event():
+        yield SSE.patch_signals({"step": 1})
+        yield SSE.patch_signals({"step": 2})
+
+    result = multi_event()
+    assert isinstance(result, DatastarResponse)
+
+    # Consume and verify events
+    events = []
+    async for event in result.body_iterator:
+        events.append(event)
+
+    assert len(events) == 2
+    assert "step" in events[0]
+    assert "step" in events[1]
+
+
+@pytest.mark.asyncio
+async def test_async_single_event():
+    """Async functions returning single event should work."""
+
+    @datastar_response
+    async def single():
+        return SSE.patch_signals({"done": True})
+
+    result = single()
+    assert isinstance(result, DatastarResponse)
+    events = []
+    async for event in result.body_iterator:
+        events.append(event)
+    assert events
+
+
+def test_sync_generator():
+    """Sync generators should yield multiple events."""
+
+    @datastar_response
+    def sync_multi():
+        yield SSE.patch_signals({"count": 1})
+        yield SSE.patch_signals({"count": 2})
+
+    result = sync_multi()
+    assert isinstance(result, DatastarResponse)
+
+
+def test_sync_single_event():
+    """Sync functions returning single event should work."""
+
+    @datastar_response
+    def single():
+        return SSE.patch_signals({"value": 42})
+
+    result = single()
+    assert isinstance(result, DatastarResponse)
+
+
+# Optional smoke tests for other adapters to ensure decorator returns the right type
+if _has("litestar"):
+    from datastar_py.litestar import DatastarResponse as LitestarResponse
+    from datastar_py.litestar import datastar_response as litestar_response
+
+    def test_litestar_sync_generator():
+        @litestar_response
+        def sync_multi():
+            yield SSE.patch_signals({"count": 1})
+            yield SSE.patch_signals({"count": 2})
+
+        resp = sync_multi()
+        assert isinstance(resp, LitestarResponse)
+
+    @pytest.mark.asyncio
+    async def test_litestar_async_function():
+        @litestar_response
+        async def single():
+            return SSE.patch_signals({"done": True})
+
+        resp = single()
+        assert isinstance(resp, LitestarResponse)
+        events = []
+        async for event in resp.iterator:
+            events.append(event)
+        assert events
+
+
+if _has("django"):
+    from datastar_py.django import DatastarResponse as DjangoResponse
+    from datastar_py.django import datastar_response as django_response
+
+    @pytest.mark.asyncio
+    async def test_django_async_function():
+        @django_response
+        async def single():
+            return SSE.patch_signals({"done": True})
+
+        resp = await single()
+        assert isinstance(resp, DjangoResponse)
+
+    @pytest.mark.asyncio
+    async def test_django_sync_generator():
+        @django_response
+        def sync_multi():
+            yield SSE.patch_signals({"count": 1})
+            yield SSE.patch_signals({"count": 2})
+
+        resp = await sync_multi()
+        assert isinstance(resp, DjangoResponse)