From c7f0c46029763f5fd8d1a8e05f27f29262d06d14 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 12 Sep 2026 10:05:46 -0700 Subject: [PATCH] refactor test --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/execution.rs | 132 +++++++++++------- 3 files changed, 83 insertions(+), 51 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 82211f5e329..1d3bf4bcbb6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1982,6 +1982,7 @@ dependencies = [ "litellm-token-counter", "pyo3", "pyo3-async-runtimes", + "rstest", "serde", "serde_json", "tokio", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 337a1e8e5ac..e432cab28dc 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -35,6 +35,7 @@ tokio = { workspace = true, features = ["sync"] } [dev-dependencies] criterion.workspace = true +rstest.workspace = true tokio-tungstenite.workspace = true tracing.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/execution.rs b/litellm-rust/crates/python-bridge/src/execution.rs index 99b1a086684..d8dda10068d 100644 --- a/litellm-rust/crates/python-bridge/src/execution.rs +++ b/litellm-rust/crates/python-bridge/src/execution.rs @@ -168,11 +168,30 @@ mod tests { use litellm_core::error::Error; use pyo3::panic::PanicException; use pyo3::types::{PyDict, PyModule}; + use rstest::{fixture, rstest}; use serde::Serializer; use tokio::runtime::Builder; use super::*; + struct InitializedPython; + + impl InitializedPython { + fn attach(&self, f: F) -> R + where + F: for<'py> FnOnce(Python<'py>) -> R, + { + Python::attach(f) + } + } + + #[fixture] + #[once] + fn initialized_python() -> InitializedPython { + Python::initialize(); + InitializedPython + } + fn runtime_error(error: Error) -> PyErr { PyRuntimeError::new_err(error.to_string()) } @@ -243,10 +262,11 @@ mod tests { .expect("result should convert") } - #[test] - fn inline_poll_releases_gil_and_enters_runtime() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn inline_poll_releases_gil_and_enters_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let (sender, receiver) = mpsc::sync_channel(1); let worker = thread::spawn(move || Python::attach(|_| sender.send(()).unwrap())); let mut future = Box::pin(async move { @@ -261,10 +281,11 @@ mod tests { }); } - #[test] - fn inline_poll_contains_panics_and_preserves_python_errors() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn inline_poll_contains_panics_and_preserves_python_errors( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let mut panicking = Box::pin(poll_fn(|_| -> Poll> { panic!("inline native panic") })); @@ -292,10 +313,11 @@ mod tests { run_async_value(py, future) } - #[test] - fn inline_pending_future_resumes_on_tokio_without_restarting() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn inline_pending_future_resumes_on_tokio_without_restarting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let locals = PyDict::new(py); locals .set_item( @@ -313,10 +335,11 @@ mod tests { }); } - #[test] - fn sync_runner_polls_future_on_the_caller_thread() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_polls_future_on_the_caller_thread( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let caller_thread = std::thread::current().id(); let result = run_sync( py, @@ -328,10 +351,11 @@ mod tests { }); } - #[test] - fn sync_runner_releases_gil_while_waiting() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_releases_gil_while_waiting( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let result = run_sync( py, async { @@ -349,16 +373,17 @@ mod tests { }); } - #[test] - fn sync_runner_rejects_calls_from_a_tokio_context() { - Python::initialize(); + #[rstest] + fn sync_runner_rejects_calls_from_a_tokio_context( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); let error = runtime.block_on(async { - Python::attach(|py| { + python.attach(|py| { run_sync::(py, async { Ok(true) }, runtime_error) .expect_err("sync route should reject a nested Tokio runtime") }) @@ -370,14 +395,15 @@ mod tests { ); } - #[test] - fn sync_runner_can_drive_a_current_thread_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_can_drive_a_current_thread_runtime( + #[from(initialized_python)] python: &InitializedPython, + ) { let runtime = Builder::new_current_thread() .enable_all() .build() .expect("runtime should build"); - Python::attach(|py| { + python.attach(|py| { let result = run_sync_on( py, &runtime, @@ -391,10 +417,9 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_future() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_future(#[from(initialized_python)] python: &InitializedPython) { + python.attach(|py| { let error = run_sync::( py, poll_fn(|_| -> Poll> { panic!("route future panicked") }), @@ -407,10 +432,11 @@ mod tests { }); } - #[test] - fn sync_runner_maps_a_panicked_error_mapper() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_maps_a_panicked_error_mapper( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync::( py, async { Err(Error::InvalidRequest("invalid".to_string())) }, @@ -423,10 +449,11 @@ mod tests { }); } - #[test] - fn sync_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn sync_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let error = run_sync(py, async { Ok(PanickingOutput) }, runtime_error) .expect_err("serializer panic should become a Python exception"); @@ -435,9 +462,10 @@ mod tests { }); } - #[test] - fn sync_runner_supports_concurrent_callers_on_the_shared_runtime() { - Python::initialize(); + #[rstest] + fn sync_runner_supports_concurrent_callers_on_the_shared_runtime( + #[from(initialized_python)] _python: &InitializedPython, + ) { let barrier = Arc::new(tokio::sync::Barrier::new(2)); let callers: Vec<_> = (0..2) .map(|_| { @@ -468,10 +496,11 @@ mod tests { assert_eq!(results, vec![true, true]); } - #[test] - fn async_runner_surfaces_serializer_panics() { - Python::initialize(); - Python::attach(|py| { + #[rstest] + fn async_runner_surfaces_serializer_panics( + #[from(initialized_python)] python: &InitializedPython, + ) { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); module .add_function( @@ -505,11 +534,12 @@ asyncio.run(exercise()) }); } - #[test] - fn async_result_delivery_does_not_stall_tokio_workers() { - Python::initialize(); + #[rstest] + fn async_result_delivery_does_not_stall_tokio_workers( + #[from(initialized_python)] python: &InitializedPython, + ) { ASYNC_PROBE_COMPLETED.store(0, Ordering::SeqCst); - Python::attach(|py| { + python.attach(|py| { let module = PyModule::new(py, "runtime").expect("module should be created"); for function in [ wrap_pyfunction!(async_runtime_probe, &module).expect("function should wrap"),