mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
ba6b22cf56
Replace the module-level LATEST_EDITS list with per-example callback registry isolation, and import litellm names with from-imports in the legacy callback shim so the module uses one import style. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
595 lines
21 KiB
Python
595 lines
21 KiB
Python
import asyncio
|
|
import copy
|
|
import gc
|
|
import queue
|
|
import threading
|
|
from collections.abc import Mapping
|
|
from types import MappingProxyType
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from hypothesis import HealthCheck, given, settings
|
|
from hypothesis import strategies as st
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
|
from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging
|
|
from tests.test_litellm_rust.support.isolation import isolated_callback_registries
|
|
from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec
|
|
from tests.test_litellm_rust.support.requests import (
|
|
OCR_DOCUMENT,
|
|
OCR_RESPONSE,
|
|
call_native,
|
|
call_native_aocr,
|
|
call_native_ocr,
|
|
request_body,
|
|
request_headers,
|
|
)
|
|
|
|
pytestmark = pytest.mark.requires_rust_extension
|
|
|
|
|
|
@pytest.fixture
|
|
def ocr_server(recording_server: RecordingServer) -> RecordingServer:
|
|
recording_server.default_response = ResponseSpec(body=OCR_RESPONSE)
|
|
return recording_server
|
|
|
|
|
|
def call_native_ocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
|
return call_native_ocr(server, callbacks=callbacks, **kwargs)
|
|
|
|
|
|
async def call_native_aocr_with_callbacks(server: RecordingServer, callbacks: list[CustomLogger], **kwargs: object):
|
|
return await call_native_aocr(server, callbacks=callbacks, **kwargs)
|
|
|
|
|
|
def test_native_ocr_pre_call_callback_receives_transformed_provider_request(ocr_server: RecordingServer) -> None:
|
|
observations: Final = []
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observations.append((model, copy.deepcopy(kwargs["additional_args"])))
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [Observe()], pages=[0])
|
|
|
|
assert len(observations) == 1
|
|
model, additional_args = observations[0]
|
|
assert model == "mistral-ocr-latest"
|
|
assert additional_args["api_base"] == f"{ocr_server.base_url}/v1/ocr"
|
|
assert additional_args["complete_input_dict"] == {
|
|
"model": "mistral-ocr-latest",
|
|
"document": OCR_DOCUMENT,
|
|
"pages": [0],
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("raise_after_edit", [False, True], ids=["callback-returns", "callback-raises"])
|
|
def test_native_ocr_pre_call_body_edit_reaches_next_callback_and_provider(
|
|
ocr_server: RecordingServer, raise_after_edit: bool
|
|
) -> None:
|
|
observed: Final = []
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs)["include_image_base64"] = True
|
|
if raise_after_edit:
|
|
raise RuntimeError("pre-call callback failed")
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(copy.deepcopy(request_body(kwargs)))
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()], include_image_base64=False)
|
|
|
|
assert observed[0]["include_image_base64"] is True
|
|
assert ocr_server.requests[0].body["include_image_base64"] is True
|
|
|
|
|
|
def test_native_ocr_pre_call_header_edit_reaches_next_callback_and_provider(ocr_server: RecordingServer) -> None:
|
|
observed: Final = []
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_headers(kwargs)["x-audit-tag"] = "reviewed"
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(dict(request_headers(kwargs)))
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [Edit(), Observe()])
|
|
|
|
assert observed[0]["x-audit-tag"] == "reviewed"
|
|
assert ocr_server.requests[0].headers["x-audit-tag"] == "reviewed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_and_provider_references(
|
|
ocr_server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
original: Final = dict(OCR_DOCUMENT)
|
|
replacement_url: Final = "data:application/pdf;base64,ZGVm"
|
|
retained: Final = []
|
|
aliases: Final = []
|
|
|
|
class Retain(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
aliases.append(request_body(kwargs)["document"] is original)
|
|
retained.append(request_body(kwargs)["document"])
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
original["document_url"] = replacement_url
|
|
|
|
arguments: Final = {
|
|
"model": "mistral/mistral-ocr-latest",
|
|
"document": original,
|
|
"api_key": "test-key",
|
|
"api_base": ocr_server.base_url,
|
|
"callbacks": [Retain(), Edit()],
|
|
}
|
|
response: Final = (
|
|
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
|
|
assert aliases == [True]
|
|
assert retained[0]["document_url"] == replacement_url
|
|
assert original["document_url"] == replacement_url
|
|
assert ocr_server.requests[0].body["document"]["document_url"] == replacement_url
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
|
|
|
|
def test_native_ocr_pre_call_body_rebinding_is_visible_to_callbacks_but_not_provider(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
observed: Final = []
|
|
|
|
class Rebind(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["additional_args"]["complete_input_dict"] = {"replacement": True}
|
|
|
|
class Observe(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
observed.append(request_body(kwargs))
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [Rebind(), Observe()])
|
|
|
|
assert observed == [{"replacement": True}]
|
|
assert ocr_server.requests[0].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT}
|
|
|
|
|
|
def test_native_ocr_callback_retained_body_observes_later_callback_mutation(ocr_server: RecordingServer) -> None:
|
|
queued: Final = []
|
|
|
|
class QueuePayload(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
queued.append(request_body(kwargs))
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs)["queued-edit"] = True
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [QueuePayload(), Edit()])
|
|
|
|
assert queued[0]["queued-edit"] is True
|
|
|
|
|
|
def test_native_ocr_success_callback_receives_state_added_by_pre_call_callback(ocr_server: RecordingServer) -> None:
|
|
token: Final = object()
|
|
terminal_tokens: queue.SimpleQueue[object] = queue.SimpleQueue()
|
|
finished: Final = threading.Event()
|
|
|
|
class Stash(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["test-token"] = token
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
terminal_tokens.put(kwargs["test-token"])
|
|
finished.set()
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [Stash()])
|
|
|
|
assert finished.wait(10)
|
|
assert terminal_tokens.get_nowait() is token
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_success_callback_receives_call_id_metadata_and_response(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
await call_native_aocr_with_callbacks(
|
|
ocr_server,
|
|
[recorder],
|
|
litellm_call_id="ocr-success",
|
|
metadata={"source": "callback-test"},
|
|
)
|
|
events: Final = await recorder.wait_for_async("async_log_success_event")
|
|
|
|
assert len(events) == 1
|
|
assert events[0].call_type == "aocr"
|
|
assert events[0].kwargs["litellm_call_id"] == "ocr-success"
|
|
assert events[0].kwargs["litellm_params"]["metadata"]["source"] == "callback-test"
|
|
assert events[0].response.pages[0].markdown == "native OCR response"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_response(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
|
|
observations: Final = []
|
|
|
|
class Observe(CustomLogger):
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observations.append(("sync", kwargs["call_type"], kwargs["exception"], response_obj))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observations.append(("async", kwargs["call_type"], kwargs["exception"], response_obj))
|
|
|
|
with pytest.raises(litellm.InternalServerError):
|
|
await call_native_aocr_with_callbacks(ocr_server, [Observe()])
|
|
|
|
assert [observation[0] for observation in observations] == ["sync", "async"]
|
|
assert all(observation[1] == "aocr" for observation in observations)
|
|
assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations)
|
|
assert all(observation[3] is None for observation in observations)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_pre_call_callback_runs_on_caller_loop_and_thread(ocr_server: RecordingServer) -> None:
|
|
caller_loop: Final = asyncio.get_running_loop()
|
|
caller_thread: Final = threading.current_thread()
|
|
recorder: Final = RecordingLogger()
|
|
|
|
await call_native_aocr_with_callbacks(ocr_server, [recorder])
|
|
|
|
events: Final = await recorder.wait_for_async("log_pre_api_call")
|
|
assert len(events) == 1
|
|
assert events[0].loop is caller_loop
|
|
assert events[0].thread is caller_thread
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_callback(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
ocr_server.enqueue(ResponseSpec(body={"message": "provider unavailable"}, status=500))
|
|
token: Final = object()
|
|
observed: Final = []
|
|
|
|
class TrackInFlightRequest(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
kwargs["request-token"] = token
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("sync", kwargs["request-token"]))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("async", kwargs["request-token"]))
|
|
|
|
with pytest.raises(litellm.InternalServerError):
|
|
await call_native_aocr_with_callbacks(ocr_server, [TrackInFlightRequest()])
|
|
|
|
assert [event for event, _ in observed] == ["sync", "async"]
|
|
assert all(observed_token is token for _, observed_token in observed)
|
|
|
|
|
|
def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registered_multiple_times(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
call_native_ocr_with_callbacks(
|
|
ocr_server,
|
|
[recorder, recorder],
|
|
success_callback=[recorder],
|
|
failure_callback=[recorder],
|
|
)
|
|
recorder.wait_for("log_success_event")
|
|
|
|
assert recorder.names.count("log_pre_api_call") == 1
|
|
assert recorder.names.count("logging_hook") == 1
|
|
assert recorder.names.count("log_success_event") == 1
|
|
assert "log_failure_event" not in recorder.names
|
|
|
|
|
|
JSON_SCALARS: Final = (
|
|
st.none()
|
|
| st.booleans()
|
|
| st.integers(min_value=-(2**63), max_value=2**63 - 1)
|
|
| st.floats(allow_nan=False, allow_infinity=False)
|
|
| st.text(max_size=8)
|
|
)
|
|
JSON_VALUES: Final = st.recursive(
|
|
JSON_SCALARS,
|
|
lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3),
|
|
max_leaves=8,
|
|
)
|
|
|
|
|
|
class ApplyEdits(CustomLogger):
|
|
def __init__(self, edits: Mapping[str, object]) -> None:
|
|
super().__init__()
|
|
self.edits: Final = edits
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
request_body(kwargs).update(copy.deepcopy(dict(self.edits)))
|
|
|
|
|
|
@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture])
|
|
@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3))
|
|
def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it(
|
|
ocr_server: RecordingServer, edits: dict[str, object]
|
|
) -> None:
|
|
ocr_server.expected_requests = None
|
|
|
|
with isolated_callback_registries():
|
|
call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))])
|
|
|
|
assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits}
|
|
|
|
|
|
@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"])
|
|
def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None:
|
|
recorder: Final = RecordingLogger()
|
|
|
|
call_native_ocr_with_callbacks(ocr_server, [recorder])
|
|
|
|
[event] = recorder.wait_for(hook)
|
|
assert event.loop is None
|
|
assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact(
|
|
ocr_server: RecordingServer, asynchronous: bool
|
|
) -> None:
|
|
retained: Final = []
|
|
|
|
class Retain(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
retained.append((kwargs, request_body(kwargs), request_headers(kwargs)))
|
|
|
|
await call_native(ocr_server, asynchronous, callbacks=[Retain()])
|
|
await drain_logging()
|
|
gc.collect()
|
|
|
|
[(details, body, headers)] = retained
|
|
assert body == ocr_server.requests[0].body
|
|
assert headers
|
|
assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items())
|
|
assert details["additional_args"]["complete_input_dict"] is body
|
|
assert details["additional_args"]["headers"] is headers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("family", ["sync", "async"])
|
|
async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None:
|
|
queued: Final = []
|
|
finished: Final = threading.Event()
|
|
|
|
def queue_payload(kwargs: dict[str, object]) -> None:
|
|
queued.append(kwargs["standard_logging_object"])
|
|
|
|
def strip_payload(kwargs: dict[str, object]) -> None:
|
|
payload: Final = kwargs["standard_logging_object"]
|
|
assert isinstance(payload, dict)
|
|
payload["stripped-by-a-later-callback"] = True
|
|
finished.set()
|
|
|
|
class QueuePayload(CustomLogger):
|
|
if family == "sync":
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
queue_payload(kwargs)
|
|
|
|
else:
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
queue_payload(kwargs)
|
|
|
|
class StripPayload(CustomLogger):
|
|
if family == "sync":
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
strip_payload(kwargs)
|
|
|
|
else:
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
strip_payload(kwargs)
|
|
|
|
await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()])
|
|
await drain_logging()
|
|
|
|
assert await asyncio.to_thread(finished.wait, 10)
|
|
assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks(
|
|
ocr_server: RecordingServer,
|
|
) -> None:
|
|
token: Final = object()
|
|
observed: Final = []
|
|
|
|
class Blocked(Exception):
|
|
pass
|
|
|
|
class Block(CustomLogger):
|
|
async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
|
|
request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token
|
|
raise Blocked("blocked after the provider answered")
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("success", None, None))
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"]))
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"]))
|
|
|
|
litellm.callbacks.append(Block())
|
|
|
|
with pytest.raises(Blocked) as raised:
|
|
await call_native_aocr(ocr_server)
|
|
await drain_logging()
|
|
|
|
assert observed == [("sync", token, raised.value), ("async", token, raised.value)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
from contextvars import ContextVar
|
|
|
|
context: Final = ContextVar("azure-token-context", default="missing")
|
|
context.set("caller")
|
|
caller_thread: Final = threading.current_thread()
|
|
caller_loop: Final = asyncio.get_running_loop()
|
|
observations: Final = []
|
|
|
|
class Provider:
|
|
def __call__(self) -> str:
|
|
assert context.get() == "caller"
|
|
assert threading.current_thread() is caller_thread
|
|
assert asyncio.get_running_loop() is caller_loop
|
|
observations.append("token")
|
|
return "caller-token"
|
|
|
|
class Edit(CustomLogger):
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
assert request_headers(kwargs)["Authorization"] == "Bearer caller-token"
|
|
observations.append("pre_call")
|
|
request_headers(kwargs)["Authorization"] = "Bearer edited"
|
|
|
|
provider: Final = Provider()
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": provider,
|
|
"callbacks": [Edit()],
|
|
}
|
|
response: Final = (
|
|
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert observations == ["token", "pre_call"]
|
|
assert ocr_server.requests[0].headers["authorization"] == "Bearer edited"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
|
async def test_native_azure_ocr_token_provider_can_make_nested_native_ocr_call(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
asynchronous: bool,
|
|
) -> None:
|
|
ocr_server.expected_requests = 2
|
|
calls: Final = []
|
|
|
|
def provider() -> str:
|
|
calls.append("token")
|
|
nested: Final = call_native_ocr(ocr_server)
|
|
assert nested.pages[0].markdown == "native OCR response"
|
|
return "outer-token"
|
|
|
|
arguments: Final = {
|
|
"model": "azure_ai/mistral-ocr-latest",
|
|
"api_key": None,
|
|
"azure_ad_token_provider": provider,
|
|
}
|
|
response: Final = (
|
|
await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments)
|
|
)
|
|
assert response.pages[0].markdown == "native OCR response"
|
|
assert calls == ["token"]
|
|
assert [request.headers["authorization"] for request in ocr_server.requests] == [
|
|
"Bearer test-key",
|
|
"Bearer outer-token",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_native_azure_ocr_calls_isolate_token_results_and_error(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
ocr_server.expected_requests = 2
|
|
|
|
async def request(token: str, fail: bool) -> object:
|
|
def provider() -> str:
|
|
if fail:
|
|
raise ValueError(token)
|
|
return token
|
|
|
|
return await call_native_aocr(
|
|
ocr_server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
|
|
responses: Final = await asyncio.gather(
|
|
request("first", False),
|
|
request("failed", True),
|
|
request("second", False),
|
|
return_exceptions=True,
|
|
)
|
|
assert isinstance(responses[0], OCRResponse)
|
|
assert isinstance(responses[1], litellm.APIConnectionError)
|
|
assert "Failed to get Azure AD token: failed" in str(responses[1])
|
|
assert isinstance(responses[2], OCRResponse)
|
|
assert sorted(request.headers["authorization"] for request in ocr_server.requests) == [
|
|
"Bearer first",
|
|
"Bearer second",
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_azure_ocr_releases_token_provider_after_cancellation(
|
|
ocr_server: RecordingServer,
|
|
isolated_azure_auth: None,
|
|
) -> None:
|
|
import gc
|
|
import weakref
|
|
|
|
from tests.test_litellm_rust.support.callback_recorder import drain_logging
|
|
|
|
class Provider:
|
|
def __call__(self) -> str:
|
|
return "caller-token"
|
|
|
|
async def invoke() -> weakref.ReferenceType[Provider]:
|
|
provider: Final = Provider()
|
|
reference: Final = weakref.ref(provider)
|
|
ocr_server.enqueue(ResponseSpec(body=OCR_RESPONSE, delay=0.1))
|
|
task: Final = asyncio.create_task(
|
|
call_native_aocr(
|
|
ocr_server,
|
|
model="azure_ai/mistral-ocr-latest",
|
|
api_key=None,
|
|
azure_ad_token_provider=provider,
|
|
)
|
|
)
|
|
await ocr_server.wait_for_requests(1)
|
|
assert reference() is provider
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
return reference
|
|
|
|
reference: Final = await invoke()
|
|
await drain_logging()
|
|
await asyncio.sleep(0)
|
|
gc.collect()
|
|
assert reference() is None
|