litellm/tests/unit/interactions/test_interactions_streaming_iterator.py
yuneng-jiang f6882246d4
test: move tests/test_litellm root and small trees into tests/unit (#43186)
* ci: run the unit_selection.sh shard files on every event instead of only fork pull requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* ci: rename fork-flag to unit-flag now that it applies on every event

* test: move tests/test_litellm root and small trees into tests/unit

Pure renames, no content changes. Follow-up commits in this PR fix
references, merge the three files that already existed in tests/unit,
keep live-provider tests in tests/test_litellm and wire CI.

* test: carry tests/test_litellm conftest isolation into tests/unit

Callback lists, routing fallbacks, cached HTTP clients, logger state, AWS,
proxy-URL and keychain env, and session-end client cleanup now reset for
unit tests too. The environment isolation owns its MonkeyPatch so a test's
own monkeypatch is undone before the model-cost teardown runs.

* test: merge, split and prune the moved root and small-tree tests

Merge batches/test_batch_utils.py and the chat_completions and messages
dispatch tests into the files that already existed in tests/unit. Keep
the live Gemini interactions tests, the async image-fetch format test and
the OpenAI embedding scorer test in tests/test_litellm since they need
real network or keys. Put test_router.py under tests/unit/test_router so
the existing package no longer shadows it. Delete eight tests the audit
found superseded by stronger ones kept in this move.

* ci: run the moved root and small-tree tests under their legacy flags

Add the misc and responses-caching-types flags to unit_selection.sh and
CircleCI, extend enterprise-routing and mcp-integration, and point the
legacy GHA shards, Makefile, redis-compat workflow, merge smoke manifest
and change classifier at the new paths.

* test: make the new tests/unit directories packages

tests/unit/test_package_layout.py requires every directory to carry an
__init__.py, and without one the moved and retained
test_litellm_responses_bridge.py modules collide on import.

* test: scope the unit socket block to tests/unit in shared sessions

The GHA shards collect the legacy test-path and the unit selection in one
pytest session. The unit conftest's loopback-only block leaked into legacy
modules that reach the network at import. The legacy conftest now lifts the
restriction at collect and setup time, and the unit conftest re-applies it
when collecting its own modules.

* test: give the shard-script tests their own GITHUB_OUTPUT

They only passed where the runner set it. The CircleCI unit job's env
allowlist drops it, so the script's redirect failed there.

* test: point the router and module-deletion checks at tests/unit

router_code_coverage and code_qa_check_tests only searched tests/test_litellm,
so the moved router tests no longer counted. The two silent-experiment tests
the audit deleted were the only direct callers of those methods; they are
replaced with tests that assert the forwarded shadow request and the
recursion guard.

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-25 11:30:43 -07:00

97 lines
3.2 KiB
Python

"""
Regression test for LIT-4210: completing an async Interactions API stream must
not run the sync success_handler on the thread-pool executor concurrently with
async_success_handler (cross-thread pydantic mutation segfaults pydantic-core).
"""
import asyncio
import time
import httpx
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.interactions import streaming_iterator as interactions_streaming_iterator_module
from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator
from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.types.interactions import InteractionsAPIStreamingResponse
class RecordingCustomLogger(CustomLogger):
def __init__(self):
super().__init__()
self.async_hook_fired = False
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.async_hook_fired = True
async def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
self.async_hook_fired = True
class RecordingExecutor:
def __init__(self, inner):
self._inner = inner
self.submits: list = []
def submit(self, fn, *args, **kwargs):
self.submits.append(fn)
return self._inner.submit(fn, *args, **kwargs)
def submitted_for(self, logging_obj) -> list:
return [fn for fn in self.submits if getattr(fn, "__self__", None) is logging_obj]
@pytest.fixture(autouse=True)
def _isolate_callbacks():
saved = (
litellm.callbacks,
litellm.success_callback,
litellm._async_success_callback,
litellm.failure_callback,
litellm._async_failure_callback,
)
yield
(
litellm.callbacks,
litellm.success_callback,
litellm._async_success_callback,
litellm.failure_callback,
litellm._async_failure_callback,
) = saved
@pytest.mark.asyncio
async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch):
recording_executor = RecordingExecutor(thread_pool_executor_module.executor)
monkeypatch.setattr(thread_pool_executor_module, "executor", recording_executor)
monkeypatch.setattr(interactions_streaming_iterator_module, "executor", recording_executor)
recorder = RecordingCustomLogger()
litellm.success_callback = [recorder]
litellm._async_success_callback = [recorder]
logging_obj = LitellmLogging(
model="gemini/gemini-3-pro-preview",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="ainteraction",
start_time=time.time(),
litellm_call_id="lit-4210-test",
function_id="lit-4210-test",
)
iterator = InteractionsAPIStreamingIterator(
response=httpx.Response(200),
model="gemini/gemini-3-pro-preview",
interactions_api_config=None,
logging_obj=logging_obj,
)
iterator.completed_response = InteractionsAPIStreamingResponse()
iterator._handle_logging_completed_response()
await asyncio.sleep(0.5)
assert recorder.async_hook_fired is True
assert recording_executor.submitted_for(logging_obj) == []