refactor test

litellm-rust/Cargo.lock

@@ -1982,6 +1982,7 @@ dependencies = [
  "litellm-token-counter",
  "pyo3",
  "pyo3-async-runtimes",
+ "rstest",
  "serde",
  "serde_json",
  "tokio",

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
 

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<F, R>(&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<PyResult<()>> {
                 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::<bool, Error, _>(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::<bool, Error, _>(
                 py,
                 poll_fn(|_| -> Poll<Result<bool, Error>> { 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::<bool, Error, _>(
                 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"),