Mocking API calls with Rust + Mockito

The main API from crivo.ai was written in Rust. One of its main purposes it is to interface with OpenAI API endpoints, mainly the embeddings and chat endpoints.

At the time I started to code it, there were only a few crates to interface Rust with the OpenAI API, and since I was willing to get better at the language, I chose to write all the interfaces on my own.

Since I may use the same codebase (regarding the OpenAI API interface) in other projects in the future, why not make it publicly available on my github? So you can check it out here.

This post is about an important step when you are planning to make your work available to others to use and contribute to.

It is about writing tests.

More specifically the testing case regarding the embeddings endpoint. Anyway, if you don't know what the embeddings are, check my post where I tried to explain about it.

The problem

I want to write a unit testing to the embeddings endpoint. In any programming language, you can find a particular feature that can help with this, and Rust provides them as well with its Test functions.

But unlike a standard synchronous function where you enter a parameter and the function returns a value, and then test it with functions like assert_eq!(value, expected_value), when you're dealing with real API calls you may not want to spend money every time you run your tests.

The solution for this is to mock the actual API call using a fake server that intercepts the request and returns a predefined response.

Mocking POST Request

If you are a Java developer probably already know about Mockito, it is a unit testing framework for Java. But when talking about Rust, mockito is a library for generating and delivering HTTP mocks in Rust.

With it, I can create a fake server to mock the HTTP Post request to test my generate_embeddings method without actually hitting the OpenAI API.

For this, it is important to know what the actual response looks like, so you can write a similar response that will be delivered by the mock server.

Below you will find the whole testing function to test when the API returns a 200 response.

use crate::OpenAi;

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_generate_embeddings_success() { // [1]
    
    let mut server = mockito::Server::new_async().await; // [2]
    let url = server.url(); // [3]
    let endpoint = format!("{}/v1/embeddings", url); // [3]

    let mock_response = r#"{ 
        "data": [
            {
                "embedding": [0.0, 0.1,0.3]
            }
        ]
    }"#; // [4]

    let mock = server.mock("POST", "/v1/embeddings")
        .with_status(200)
        .with_header("content-type", "application/json")
        .with_body(mock_response)
        .create_async().await; // [5]

    
    let client = OpenAi::new("you_look_lonely_I_can_fix_that".to_string())
        .set_endpoint(endpoint); // [6]

    let result = client.generate_embeddings("I am lonely", "some-wonderful-embedding-model").await; // [7]

    assert!(result.is_ok()); // [8]

    let embeddings_vec = result.unwrap();

    assert_eq!(embeddings_vec, vec![0.0, 0.1, 0.3]); // [8]

    mock.assert_async().await; // [9]

}

I will break down the important steps of the code below.

1. Since we want to make actual requests to a third-party API it is important to make sure your function is asynchronous. The "almost official" way to deal with async functions in Rust is by the Tokyo crate.
2. The mockito server is instantiated as async as well since it will be waiting for a further call.
3. The URL from the fake server is obtained and the embedding endpoint is merged to it.
4. The mock response is then defined according to the actual expected response.
5. The mock server is configured to listen for POST requests to /v1/embeddings and to respond with status 200 with the fake response defined in step 4.
6. A new instance of the OpenAI client is created with a dummy API KEY and sets the endpoint to the mock server.
7. The method responsible for interacting with OpenAI embeddings endpoint is then called according to its signature and the request is sent to the mock server.
8. The official Rust asset functions are called to check for STATUS 200 and to check if the generate_embeddings method is returning the embeddings vector as expected.
9. The last assertion checks if the mock endpoint was called during the test.

Conclusion

Writing tests is a fundamental step to ensure the reliability and correctness of your code, especially when you plan to share it with others or reuse it in future projects. By mocking API requests using tools like mockito, you can simulate real-world scenarios without spending money on real calls to external APIs.

If you haven’t already, consider exploring the full example and trying it out yourself. Feel free to reach out with feedback or share how you implemented HTTP mock tests on your projects.