litellm/litellm-rust/crates/host-python/tests/driver.rs
yujonglee 327bdcaa7f
test(host-python): move driver tests out of the inline module (#46325)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-14 17:19:05 -07:00

866 lines
27 KiB
Rust

mod support;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_host::{
hooks::CallHooks,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
protocol::Reply,
};
use litellm_host_python::{
CallOptions, HookChain, InvokeError, PythonBinding, PythonHostCalls, PythonOwned,
PythonRuntime, missing_state, run_call,
};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
PyValueError, asyncio::CancelledError,
},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use rstest::rstest;
use support::{
Classified, Error, HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding,
call_options, failing_machine, hooks, interpreter, lifecycle_binding, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
let options = CallOptions {
asynchronous: false,
lifecycle: lifecycle_binding,
observers: Some(sender),
};
(options, move || {
std::iter::from_fn(|| receiver.try_recv().ok()).collect()
})
}
#[rstest]
#[case::success(OpScript::AwaitAnswer)]
#[case::failure(OpScript::AwaitFailure)]
fn asynchronous_host_operations_resume_the_same_machine(
#[from(interpreter)] _python: Interpreter,
#[case] op: OpScript,
) {
Python::attach(|py| {
let machine: MachineFactory = Box::new(|_| {
CallMachine::<Synthetic>::new(None, |host| {
Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
})
});
let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
match op {
OpScript::AwaitAnswer => {
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
}
OpScript::AwaitFailure => {
assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
1
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
}
_ => unreachable!(),
}
});
}
#[rstest]
fn composed_hooks_share_prepared_arguments_and_one_execution(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let chain = HookChain::new()
.with(hooks(&log, HookScript::ObserveArguments))
.with(hooks(&log, HookScript::RewriteArguments))
.with(hooks(&log, HookScript::ReplaceResponse));
let result = run_call(
py,
move |_, _, request| Ok(success_machine()(request)),
binding(&log, OpScript::Answer),
chain,
PyDict::new(py).unbind(),
call_options(asynchronous),
);
if asynchronous {
assert!(log.entries().is_empty());
}
let response = settled(py, result, asynchronous).unwrap();
assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
let entries = log.entries();
assert!(entries.iter().any(|entry| entry == "adopted:hook"));
assert_eq!(
entries.iter().filter(|entry| *entry == "project").count(),
1
);
assert_eq!(
entries.iter().filter(|entry| *entry == "op:sign").count(),
1
);
assert_eq!(
entries
.iter()
.filter(|entry| *entry == "succeeded:replaced")
.count(),
3
);
assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
});
}
#[rstest]
#[case::draining(4, false)]
#[case::full(1, false)]
#[case::closed(4, true)]
fn observation_delivery_preserves_callback_order_and_response(
#[from(interpreter)] _python: Interpreter,
#[case] capacity: usize,
#[case] closed: bool,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (expected, expected_log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
if closed {
receiver.close();
}
let log = Log::default();
let (actual, actual_log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::ReplaceResponse),
std::convert::identity,
CallOptions {
asynchronous,
lifecycle: lifecycle_binding,
observers: Some(sender.clone()),
},
);
assert_eq!(
actual.unwrap().extract::<String>(py).unwrap(),
expected.unwrap().extract::<String>(py).unwrap()
);
assert_eq!(actual_log, expected_log);
let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
if !closed && capacity >= 3 {
let [
CallEvent::Started { start_time },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
CallEvent::Succeeded { timing, .. },
] = events.as_slice()
else {
panic!("expected start, provider response and success: {events:?}");
};
assert_eq!(raw.body, "raw");
assert_eq!(*start_time, timing.start_time);
assert!(timing.end_time >= timing.start_time);
}
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, false, 1)]
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] cancelled: bool,
#[case] failure_callbacks: usize,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (options, mut events) = observed(4);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, script),
std::convert::identity,
CallOptions {
asynchronous,
..options
},
);
let error = result.unwrap_err();
assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
let events = events();
assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
assert_eq!(
matches!(events.last(), Some(CallEvent::Cancelled { .. })),
cancelled
);
assert_eq!(
matches!(
events.last(),
Some(CallEvent::Failed {
origin: FailureOrigin::Host,
..
})
),
!cancelled
);
let terminal = events.iter().filter(|event| {
matches!(
event,
CallEvent::Failed { .. }
| CallEvent::Succeeded { .. }
| CallEvent::Cancelled { .. }
)
});
assert_eq!(terminal.count(), 1);
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
failure_callbacks
);
assert!(
!log.iter()
.any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
);
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
fn startup_failures_release_the_factory_without_constructing_a_machine(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] op: OpScript,
#[values(false, true)] asynchronous: bool,
) {
struct Release(Arc<AtomicBool>);
impl Drop for Release {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
Python::attach(|py| {
let constructed = Arc::new(AtomicBool::new(false));
let did_construct = constructed.clone();
let released = Arc::new(AtomicBool::new(false));
let release = Release(released.clone());
let machine: MachineFactory = Box::new(move |request| {
let _release = release;
did_construct.store(true, Ordering::SeqCst);
success_machine()(request)
});
let (result, log) = run_scripted(py, machine, op, script, asynchronous);
assert!(result.is_err());
assert!(!constructed.load(Ordering::SeqCst));
assert!(released.load(Ordering::SeqCst));
assert_eq!(
log.iter()
.filter(|event| event.starts_with("failed:"))
.count(),
1
);
assert_eq!(
log.iter().filter(|event| *event == "adapter.close").count(),
1
);
assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
});
}
#[rstest]
fn success_runs_every_step_in_order_and_returns_the_public_response(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:1|sign|rewritten"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"before_provider_request",
"response:raw",
"complete",
"after_success",
"succeeded:project:1|sign|rewritten",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (options, mut events) = observed(3);
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
CallOptions {
asynchronous,
..options
},
);
if asynchronous {
for method in ["__anext__", "aclose", "aclose"] {
let stop = stream
.call_method0(method)
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
assert!(stop.is_instance_of::<PyStopIteration>(py));
}
} else {
assert_eq!(
stream
.call_method0("__next__")
.unwrap()
.extract::<String>()
.unwrap(),
"first"
);
stream.call_method0("close").unwrap();
stream.call_method0("close").unwrap();
}
assert!(matches!(
events().as_slice(),
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
));
assert_eq!(
log.entries(),
[
"started",
"begin",
"opened",
"delivered",
"succeeded:None",
"adapter.close"
]
);
});
}
#[rstest]
fn a_stream_carries_its_head_before_the_first_chunk(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
call_options(asynchronous),
);
let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
stream.getattr("head").unwrap().extract().unwrap();
assert_eq!(
hidden["additional_headers"],
std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
);
assert_eq!(log.entries(), ["started", "begin", "opened"]);
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
});
}
#[rstest]
fn a_native_failure_is_classified_once_and_reported_classified(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
failing_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert_eq!(error.value(py).to_string(), "classified: provider exploded");
assert_eq!(
log,
[
"started",
"begin",
"project",
"classify:provider exploded",
"failed:Call:classified: provider exploded",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_native_rejection_from_a_host_operation_is_classified_once(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::RejectNatively,
HookScript::Plain,
false,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"classified: op rejected"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"classify:op rejected",
"failed:Call:classified: op rejected",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_python_exception_from_a_host_operation_is_reported_as_raised(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::RaisePython,
HookScript::Plain,
false,
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyValueError>(py));
assert_eq!(error.value(py).to_string(), "op failed");
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"failed:Call:op failed",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn a_failing_classifier_surfaces_with_the_native_error_as_context(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log_handle = Log::default();
let (result, log) = run_composed(
py,
failing_machine(),
binding(&log_handle, OpScript::Answer).with_failing_classifier(),
hooks(&log_handle, HookScript::Plain),
std::convert::identity,
call_options(false),
);
let error = result.unwrap_err();
assert!(error.is_instance_of::<PyTypeError>(py));
assert_eq!(error.value(py).to_string(), "classifier failed");
let context = error.context(py).unwrap();
assert!(context.is_instance_of::<PyRuntimeError>(py));
assert_eq!(context.value(py).to_string(), "provider exploded");
assert_eq!(
log,
[
"started",
"begin",
"project",
"classify:provider exploded",
"failed:Call:classifier failed",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn begin_failures_are_host_failures_without_provider_mapping(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailBegin,
false,
);
assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:begin failed",
"adapter.close",
"host.close"
]
);
});
}
enum ArgumentPolicy {
Inherit,
Reject(Py<PyBaseException>),
}
impl CallHooks<PythonRuntime> for ArgumentPolicy {
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
match self {
Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
}
}
}
impl PythonOwned for ArgumentPolicy {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::Inherit => Ok(()),
Self::Reject(error) => visit.call(error),
}
}
}
#[rstest]
fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[values(false, true)] fail_setup: bool,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let error = PyValueError::new_err("setup failed").into_value(py);
let setup_error = error.clone_ref(py);
let arguments = PyDict::new(py);
arguments.set_item("api_key", "original").unwrap();
let result = run_call(
py,
move |py, prepared, request| {
let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
setup_log.push(format!("setup:{source}:{key}"));
if fail_setup {
return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
}
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
HookChain::new()
.with(hooks(&log, HookScript::RewriteArguments))
.with(ArgumentPolicy::Inherit),
arguments.clone().unbind(),
call_options(asynchronous),
);
if asynchronous {
assert!(log.entries().is_empty());
}
let settled = settled(py, result, asynchronous);
assert_eq!(
arguments
.get_item("api_key")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"original"
);
assert_eq!(
&log.entries()[..4],
["started", "begin", "project", "setup:hook:inherited"]
);
if fail_setup {
assert!(settled.unwrap_err().value(py).is(error.bind(py)));
assert_eq!(
log.entries(),
[
"started",
"begin",
"project",
"setup:hook:inherited",
"failed:Host:setup failed",
"adapter.close",
"host.close",
]
);
} else {
assert!(settled.is_ok());
assert!(
log.entries()
.iter()
.any(|entry| entry.starts_with("succeeded:"))
);
}
});
}
#[rstest]
fn closing_an_unstarted_call_never_initializes_resources(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let pending = run_call(
py,
move |_, _, request| {
setup_log.push("setup");
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
PyDict::new(py).unbind(),
call_options(true),
)
.unwrap();
assert!(log.entries().is_empty());
pending.call_method0(py, "close").unwrap();
drop(pending);
assert_eq!(log.entries(), ["adapter.close", "host.close"]);
});
}
#[rstest]
fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let raised = PyValueError::new_err("over budget").into_value(py);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| {
HookChain::new()
.with(hooks)
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
},
call_options(asynchronous),
);
assert!(result.unwrap_err().value(py).is(raised.bind(py)));
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:over budget",
"adapter.close",
"host.close"
]
);
});
}
#[rstest]
fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (result, _) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
call_options(asynchronous),
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:2|sign|rewritten"
);
});
}
#[rstest]
fn the_adapters_finalized_response_is_what_the_call_returns_and_reports(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "replaced");
assert!(log.contains(&"succeeded:replaced".to_string()));
assert!(!log.contains(&"succeeded:project:1|rewritten".to_string()));
});
}
#[rstest]
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailAfterSuccess,
asynchronous,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"after_success failed"
);
assert_eq!(
&log[log.len() - 4..],
[
"after_success",
"failed:Host:after_success failed",
"adapter.close",
"host.close"
]
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
});
}
#[rstest]
fn cancellation_ends_the_call_without_terminal_dispatch(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
struct Cancelling(Log);
impl PythonBinding for Cancelling {
type Protocol = Synthetic;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
_: &Bound<'_, PyDict>,
) -> Result<String, InvokeError<Error>> {
self.0.push("project");
Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
}
fn encode_stream_head(
&mut self,
_: Python<'_>,
head: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match head {}
}
fn encode_chunk(
&mut self,
_: Python<'_>,
chunk: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match chunk {}
}
fn encode_response(&mut self, _: Python<'_>, _: String) -> PyResult<Py<PyAny>> {
Err(missing_state())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
self.0.push("classify");
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Synthetic> for Cancelling {
fn handle_host_call(
&mut self,
_: Python<'_>,
_: (&'static str, Reply<String>),
) -> Result<(), InvokeError<Error>> {
Err(missing_state().into())
}
}
impl PythonOwned for Cancelling {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
let log = Log::default();
let host = Cancelling(log.clone());
let adapter = hooks(&log, HookScript::Plain);
let error = run_call(
py,
|_, _, request| Ok(success_machine()(request)),
host,
adapter,
PyDict::new(py).unbind(),
call_options(false),
)
.unwrap_err();
assert!(!error.is_instance_of::<PyException>(py));
assert_eq!(
log.entries(),
["started", "begin", "project", "adapter.close"]
);
});
}