How can I test rust functions wrapped by pyo3 in rust, before bu

ghz 昨天 ⋅ 3 views

How can I test rust functions wrapped by pyo3 in rust, before building and installing in python?

I have functionality written in rust, which I am exposing to python via pyo3. I would like to test that the python functions are correctly exposed and handle python types correctly.

I already have tests in place to validate the actual functional implementation (in rust) and the end-to-end integration (in python).

How can I test the pyo3 python functions in rust?

Answer

Testing PyO3 Python functions from within Rust can be achieved by creating Rust unit tests that directly interact with Python through the PyO3 framework. This allows you to test the Python API you have exposed using PyO3, ensuring that Python types are correctly handled and the integration works as expected.

Here’s how you can approach testing Python functions exposed via PyO3 directly in Rust:

1. Set Up a #[cfg(test)] Module in Rust

You can use PyO3's pyo3::types::Py and pyo3::Python to interact with Python objects and execute Python code from within Rust tests.

Example: Testing Python Functions Exposed by PyO3

Rust Side (PyO3 code)

Let’s assume you have a simple function exposed to Python that you want to test from Rust. Here’s a basic example of exposing a function to Python using PyO3:

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

/// A simple function that we want to expose to Python.
#[pyfunction]
fn add(a: usize, b: usize) -> usize {
    a + b
}

/// Create a Python module exposing the `add` function.
#[pymodule]
fn mymodule(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(add, py)?)?;
    Ok(())
}

Unit Test in Rust with PyO3

In Rust, you can test the Python API by calling the Python functions exposed through PyO3 in a test module:

#[cfg(test)]
mod tests {
    use super::*;
    use pyo3::types::PyString;
    use pyo3::types::PyDict;
    use pyo3::Python;

    /// A test to validate the `add` function exposed to Python.
    #[test]
    fn test_add_function() {
        // Initialize Python interpreter
        let gil = Python::acquire_gil();
        let py = gil.python();

        // Import the Python module
        let mymodule = PyModule::new(py, "mymodule").unwrap();
        
        // Add the function to the module
        mymodule.add_function(wrap_pyfunction!(add, py).unwrap()).unwrap();

        // Call the Python function directly using the module
        let result: usize = mymodule.call1("add", (5, 3)).unwrap().extract(py).unwrap();

        // Assert that the result is correct
        assert_eq!(result, 8);
    }

    /// Another test case with Python objects (testing string handling).
    #[test]
    fn test_python_string() {
        let gil = Python::acquire_gil();
        let py = gil.python();

        let result: String = PyString::new(py, "Hello, PyO3!")
            .to_object(py)
            .extract(py)
            .unwrap();

        assert_eq!(result, "Hello, PyO3!");
    }
}

Key Points in the Test Code:

  1. Acquire Python GIL: You need to acquire the Python Global Interpreter Lock (GIL) when interacting with Python objects. This is done using Python::acquire_gil().

  2. Create Python Module and Functions: In the test_add_function() example, we use PyModule::new(py, "mymodule") to create a module in Python and then add functions to that module using mymodule.add_function(wrap_pyfunction!(add, py)?).

  3. Calling Python Functions: After the function is added to the module, you can call it using call1() (passing the arguments as a tuple) and then extract the result back into Rust using extract(py).

  4. Testing Python Objects: The test_python_string() shows how to create a Python str and extract it back into a Rust string.

  5. Assertions: After the function is called and the result is extracted, use standard Rust assertions (assert_eq!) to verify the correctness of the result.

2. Running the Tests

To run the tests, you can use cargo test as usual. PyO3 will set up the Python interpreter when tests are run, and the tests will interact with the Python code directly.

cargo test

3. Why Test Python Functions in Rust?

Testing the Python functions in Rust serves several purposes:

  • Unit Test the Python Interface: Ensure that the Python interface behaves as expected when invoked from Rust.
  • Catch Type Mismatches: Ensure that types like str, int, list, etc., are correctly handled and mapped between Rust and Python.
  • Check for Proper Error Handling: Verify that Python exceptions are correctly handled by your Rust code (i.e., PyResult is used properly).

4. Considerations

  • Testing Environment: Ensure that Python is available and correctly configured on your system. PyO3 relies on the Python interpreter being available.
  • PyO3 and Python Version Compatibility: Ensure that the version of PyO3 you're using is compatible with the Python version you're testing against.
  • Complex Types: For more complex objects (e.g., custom classes), ensure that you also test how those objects are serialized/deserialized correctly between Rust and Python.

By integrating these tests into your CI/CD pipeline, you can ensure that your Python functions behave as expected and that they work seamlessly with your Rust code.

#py