Skip to content

Commit d0f14ae

Browse files
authored
test: recategorize llm based astream_incremental & add versions with mocks (#567)
* test: isolate astream_incremental tests from CI Fixes #562 * test: add deterministic mock tests for astream incremental logic Introduces `test_astream_mock.py` to test `ModelOutputThunk`'s async queue incremental streaming logic deterministically without relying on highly-variable LLM backends. * test: adapt astream mock tests to upstream incremental semantics Update tests to match the astream() behavior change from PR #618: - astream() now always returns incremental content (including final call) - astream() on a computed MOT raises RuntimeError * chore: trigger CI rebuild
1 parent ea8d21e commit d0f14ae

2 files changed

Lines changed: 169 additions & 1 deletion

File tree

test/core/test_astream_incremental.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
@pytest.mark.ollama
1616
@pytest.mark.llm
17+
@pytest.mark.qualitative
1718
async def test_astream_returns_incremental_chunks():
1819
"""Test that astream() returns only new content, not accumulated content.
1920
@@ -68,6 +69,7 @@ async def test_astream_returns_incremental_chunks():
6869

6970
@pytest.mark.ollama
7071
@pytest.mark.llm
72+
@pytest.mark.qualitative
7173
async def test_astream_multiple_calls_accumulate_correctly():
7274
"""Test that multiple astream() calls accumulate to the final value.
7375
@@ -111,6 +113,7 @@ async def test_astream_multiple_calls_accumulate_correctly():
111113

112114
@pytest.mark.ollama
113115
@pytest.mark.llm
116+
@pytest.mark.qualitative
114117
async def test_astream_beginning_length_tracking():
115118
"""Test that beginning_length is correctly tracked across astream calls.
116119
@@ -141,6 +144,7 @@ async def test_astream_beginning_length_tracking():
141144

142145
@pytest.mark.ollama
143146
@pytest.mark.llm
147+
@pytest.mark.qualitative
144148
async def test_astream_empty_beginning():
145149
"""Test astream when _underlying_value starts as None."""
146150
session = start_session()
@@ -167,7 +171,6 @@ async def test_astream_empty_beginning():
167171
@pytest.mark.llm
168172
async def test_computed_mot_raises_error_for_astream():
169173
"""Test that computed mot raises an error for astream() calls."""
170-
171174
# Create a pre-computed thunk
172175
mot = ModelOutputThunk(value="Hello, world!")
173176
mot._computed = True

test/core/test_astream_mock.py

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
"""Deterministic Mock Tests for ModelOutputThunk.astream() incremental return behavior.
2+
3+
Tests that astream() returns only the incremental content added during each call.
4+
All astream() chunks concatenated should equal the full final value. Calling
5+
astream() on a computed MOT raises RuntimeError. Uses manual queue injection to
6+
bypass LLM calls and network operations, guaranteeing determinism.
7+
"""
8+
9+
import asyncio
10+
from typing import Any
11+
12+
import pytest
13+
14+
from mellea.core.base import CBlock, GenerateType, ModelOutputThunk
15+
16+
17+
async def mock_process(mot: ModelOutputThunk, chunk: Any) -> None:
18+
"""Mock process function that simply appends the chunk to the underlying value."""
19+
if mot._underlying_value is None:
20+
mot._underlying_value = ""
21+
if chunk is not None:
22+
mot._underlying_value += chunk
23+
24+
25+
async def mock_post_process(mot: ModelOutputThunk) -> None:
26+
"""Mock post-process function (does nothing)."""
27+
28+
29+
def create_manual_mock_thunk() -> ModelOutputThunk:
30+
"""Helper to create a mock ModelOutputThunk where we manually populate the queue."""
31+
mot = ModelOutputThunk(value=None)
32+
mot._action = CBlock("mock_action")
33+
mot._generate_type = GenerateType.ASYNC
34+
mot._process = mock_process
35+
mot._post_process = mock_post_process
36+
mot._chunk_size = 0 # Read exactly what is available
37+
return mot
38+
39+
40+
@pytest.mark.asyncio
41+
async def test_astream_returns_incremental_chunks():
42+
"""Test that astream() returns only new content, not accumulated content."""
43+
mot = create_manual_mock_thunk()
44+
45+
# Drop the first chunk and pull it
46+
mot._async_queue.put_nowait("chunk1 ")
47+
chunk1 = await mot.astream()
48+
assert chunk1 == "chunk1 "
49+
50+
# Drop the second chunk and pull it
51+
mot._async_queue.put_nowait("chunk2 ")
52+
chunk2 = await mot.astream()
53+
assert chunk2 == "chunk2 "
54+
55+
# Drop the third chunk and pull it
56+
mot._async_queue.put_nowait("chunk3 ")
57+
chunk3 = await mot.astream()
58+
assert chunk3 == "chunk3 "
59+
60+
# Send completion sentinel
61+
mot._async_queue.put_nowait(None)
62+
63+
# Wait until fully consumed
64+
while not mot.is_computed():
65+
await mot.astream()
66+
67+
final_val = await mot.avalue()
68+
assert final_val == "chunk1 chunk2 chunk3 "
69+
70+
71+
@pytest.mark.asyncio
72+
async def test_astream_multiple_calls_accumulate_correctly():
73+
"""Test that multiple astream() calls accumulate to the final value."""
74+
# Simulating a scenario where queue chunks outpace the reading loop
75+
mot = create_manual_mock_thunk()
76+
77+
# Drop multiple items at once to simulate fast network
78+
mot._async_queue.put_nowait("c")
79+
mot._async_queue.put_nowait("h")
80+
mot._async_queue.put_nowait("u")
81+
82+
# Calling astream should drain all currently queued items ("chu")
83+
chunk1 = await mot.astream()
84+
assert chunk1 == "chu"
85+
86+
mot._async_queue.put_nowait("n")
87+
mot._async_queue.put_nowait("k")
88+
mot._async_queue.put_nowait(None)
89+
90+
chunk2 = await mot.astream()
91+
# astream() returns only the incremental content added during this call
92+
assert chunk2 == "nk"
93+
94+
assert mot.is_computed()
95+
# All astream() chunks concatenated should equal the full value
96+
assert chunk1 + chunk2 == "chunk"
97+
assert mot.value == "chunk"
98+
99+
100+
@pytest.mark.asyncio
101+
async def test_astream_beginning_length_tracking():
102+
"""Test that beginning_length is correctly tracked across astream calls."""
103+
mot = create_manual_mock_thunk()
104+
105+
mot._async_queue.put_nowait("AAA")
106+
chunk1 = await mot.astream()
107+
assert chunk1 == "AAA"
108+
109+
mot._async_queue.put_nowait("BBB")
110+
chunk2 = await mot.astream()
111+
# verify incremental length tracking works
112+
assert not chunk2.startswith(chunk1)
113+
assert chunk2 == "BBB"
114+
115+
116+
@pytest.mark.asyncio
117+
async def test_astream_empty_beginning():
118+
"""Test astream when _underlying_value starts as None."""
119+
mot = create_manual_mock_thunk()
120+
121+
mot._async_queue.put_nowait("First")
122+
# At the start, _underlying_value is None, beginning_length is 0
123+
chunk = await mot.astream()
124+
125+
# Because beginning length was 0, astream returns the full chunk
126+
assert chunk == "First"
127+
assert mot._underlying_value == "First"
128+
129+
130+
@pytest.mark.asyncio
131+
async def test_astream_computed_raises_error():
132+
"""Test that astream raises RuntimeError when already computed."""
133+
# Precomputed thunk is already computed
134+
mot = ModelOutputThunk(value="Hello, world!")
135+
136+
# astream() on a computed MOT now raises RuntimeError
137+
with pytest.raises(RuntimeError, match="Streaming has finished"):
138+
await mot.astream()
139+
140+
141+
@pytest.mark.asyncio
142+
async def test_astream_final_call_returns_full_value():
143+
"""Test that the final astream call returns the full value when computed."""
144+
mot = create_manual_mock_thunk()
145+
146+
mot._async_queue.put_nowait("part1")
147+
chunk1 = await mot.astream()
148+
assert chunk1 == "part1"
149+
150+
mot._async_queue.put_nowait("part2")
151+
chunk2 = await mot.astream()
152+
assert chunk2 == "part2"
153+
154+
mot._async_queue.put_nowait("part3")
155+
mot._async_queue.put_nowait(None)
156+
157+
# Calling astream here processes "part3" and `None`, flagging it as done
158+
chunk3 = await mot.astream()
159+
160+
# The final astream() call returns only the incremental content, not the full value
161+
assert chunk3 == "part3"
162+
163+
# All chunks concatenated equal the full value
164+
assert chunk1 + chunk2 + chunk3 == "part1part2part3"
165+
assert mot.value == "part1part2part3"

0 commit comments

Comments
 (0)