tonic-mock 0.4.0

use prost::Message;
use std::fmt::Debug;
use tonic::{Code, Request, Status};
use tonic_mock::client_mock::{GrpcClientExt, MockResponseDefinition, MockableGrpcClient};

// A typical gRPC client generated by tonic
#[derive(Debug, Clone)]
pub struct UserServiceClient<T> {
    inner: T,
}

// Messages for our User service
#[derive(Clone, PartialEq, Message)]
pub struct GetUserRequest {
    #[prost(string, tag = "1")]
    pub user_id: String,
}

#[derive(Clone, PartialEq, Message)]
pub struct User {
    #[prost(string, tag = "1")]
    pub id: String,
    #[prost(string, tag = "2")]
    pub name: String,
    #[prost(string, tag = "3")]
    pub email: String,
}

// Implementation of our extension trait for mocking
impl GrpcClientExt<UserServiceClient<MockableGrpcClient>>
    for UserServiceClient<MockableGrpcClient>
{
    fn with_mock(mock: MockableGrpcClient) -> Self {
        Self { inner: mock }
    }
}

// Simple implementation for our mock test - not a full tonic client implementation
impl UserServiceClient<MockableGrpcClient> {
    pub async fn get_user(
        &mut self,
        request: Request<GetUserRequest>,
    ) -> Result<tonic::Response<User>, Status> {
        // Extract the inner client
        let client = &self.inner;

        // Get the request data
        let request_data = request.into_inner();

        // For debugging
        println!("Debug: Sending request with user_id: {:?}", request_data);

        // Encode the request
        let encoded = tonic_mock::grpc_mock::encode_grpc_request(request_data);

        // Call the mock service
        let (response_bytes, http_metadata) = client
            .handle_request("user.UserService", "GetUser", &encoded)
            .await?;

        // Decode the response
        let user: User = tonic_mock::grpc_mock::decode_grpc_message(&response_bytes)?;

        // Create a tonic Response
        let mut tonic_response = tonic::Response::new(user);

        // Manually add specific metadata that might be available
        // We can't convert directly due to lifetime issues
        if let Some(key) = http_metadata.get("x-request-id") {
            if let Ok(val_str) = key.to_str() {
                tonic_response
                    .metadata_mut()
                    .insert("x-request-id", val_str.parse().unwrap());
            }
        }

        Ok(tonic_response)
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Mockable gRPC Client Example ===\n");

    println!("Example 1: Basic Usage");
    basic_usage_example().await?;

    println!("\nExample 2: Single Response");
    single_response_example().await?;

    println!("\nExample 3: Conditional Responses");
    conditional_response_example().await?;

    println!("\nExample 4: Error Handling");
    error_handling_example().await?;

    println!("\nExample 5: Simulating Delays");
    delay_example().await?;

    Ok(())
}

async fn basic_usage_example() -> Result<(), Box<dyn std::error::Error>> {
    // Create a mock client
    let mock = MockableGrpcClient::new();

    // Configure a mock response for the GetUser method
    mock.mock::<GetUserRequest, User>("user.UserService", "GetUser")
        .respond_with(
            MockResponseDefinition::ok(User {
                id: "user123".to_string(),
                name: "Alice Smith".to_string(),
                email: "alice@example.com".to_string(),
            })
            .with_metadata("x-request-id", "req-123"),
        )
        .await;

    // Create a client using the mock
    let mut client = UserServiceClient::with_mock(mock);

    // Make a request
    let request = GetUserRequest {
        user_id: "user123".to_string(),
    };

    // Call the service
    let response = client.get_user(Request::new(request)).await?;
    let user = response.get_ref();

    // Print the result
    println!("Retrieved user: {} <{}>", user.name, user.email);
    println!("Response metadata: {:?}", response.metadata());

    Ok(())
}

async fn single_response_example() -> Result<(), Box<dyn std::error::Error>> {
    // Create a mock client
    let mock = MockableGrpcClient::new();

    // Configure a response for any request
    mock.mock::<GetUserRequest, User>("user.UserService", "GetUser")
        .respond_with(MockResponseDefinition::ok(User {
            id: "default-user".to_string(),
            name: "Default User".to_string(),
            email: "default@example.com".to_string(),
        }))
        .await;

    // Create a client using the mock
    let mut client = UserServiceClient::with_mock(mock);

    // Make a request - any ID will return the same response
    let request = GetUserRequest {
        user_id: "any-id".to_string(),
    };

    // Call the service
    let response = client.get_user(Request::new(request)).await?;
    let user = response.get_ref();

    // Print the result
    println!("Retrieved user: {} <{}>", user.name, user.email);

    Ok(())
}

async fn conditional_response_example() -> Result<(), Box<dyn std::error::Error>> {
    // Create a mock client
    let mock = MockableGrpcClient::new();

    // Define our users
    let alice = User {
        id: "user123".to_string(),
        name: "Alice Smith".to_string(),
        email: "alice@example.com".to_string(),
    };

    let bob = User {
        id: "user456".to_string(),
        name: "Bob Jones".to_string(),
        email: "bob@example.com".to_string(),
    };

    // Configure different responses for each request - order matters!
    // The most specific handlers should be added first
    mock.mock::<GetUserRequest, User>("user.UserService", "GetUser")
        .respond_when(
            |req| req.user_id == "user123",
            MockResponseDefinition::ok(alice),
        )
        .await
        .respond_when(
            |req| req.user_id == "user456",
            MockResponseDefinition::ok(bob),
        )
        .await;

    // Create a client using the mock
    let mut client = UserServiceClient::with_mock(mock);

    // Make requests for different users
    for user_id in &["user123", "user456"] {
        let request = GetUserRequest {
            user_id: user_id.to_string(),
        };

        // Call the service
        let response = client.get_user(Request::new(request)).await?;
        let user = response.get_ref();

        // Print the result
        println!("User {}: {} <{}>", user_id, user.name, user.email);
    }

    Ok(())
}

async fn error_handling_example() -> Result<(), Box<dyn std::error::Error>> {
    // Create a mock client
    let mock = MockableGrpcClient::new();

    // Configure a mix of success and error responses
    mock.mock::<GetUserRequest, User>("user.UserService", "GetUser")
        .respond_when(
            |req| req.user_id == "valid-user",
            MockResponseDefinition::ok(User {
                id: "valid-user".to_string(),
                name: "Valid User".to_string(),
                email: "valid@example.com".to_string(),
            }),
        )
        .await
        .respond_when(
            |req| req.user_id == "not-found",
            MockResponseDefinition::err(Status::new(Code::NotFound, "User not found")),
        )
        .await
        .respond_when(
            |req| req.user_id == "unauthorized",
            MockResponseDefinition::err(Status::new(
                Code::PermissionDenied,
                "Not authorized to access this user",
            )),
        )
        .await;

    // Create a client using the mock
    let mut client = UserServiceClient::with_mock(mock);

    // Make requests with different outcomes
    for user_id in &["valid-user", "not-found", "unauthorized"] {
        let request = GetUserRequest {
            user_id: user_id.to_string(),
        };

        // Call the service
        match client.get_user(Request::new(request)).await {
            Ok(response) => {
                let user = response.get_ref();
                println!("Success: User {} found: {}", user.id, user.name);
            }
            Err(status) => {
                println!(
                    "Error for {}: {} (code: {:?})",
                    user_id,
                    status.message(),
                    status.code()
                );
            }
        }
    }

    Ok(())
}

async fn delay_example() -> Result<(), Box<dyn std::error::Error>> {
    // Create a mock client
    let mock = MockableGrpcClient::new();

    // Configure responses with different delays
    mock.mock::<GetUserRequest, User>("user.UserService", "GetUser")
        .respond_when(
            |req| req.user_id == "fast-user",
            MockResponseDefinition::ok(User {
                id: "fast-user".to_string(),
                name: "Fast Response".to_string(),
                email: "fast@example.com".to_string(),
            })
            .with_delay(100), // 100ms delay
        )
        .await
        .respond_when(
            |req| req.user_id == "slow-user",
            MockResponseDefinition::ok(User {
                id: "slow-user".to_string(),
                name: "Slow Response".to_string(),
                email: "slow@example.com".to_string(),
            })
            .with_delay(500), // 500ms delay
        )
        .await;

    // Create a client using the mock
    let mut client = UserServiceClient::with_mock(mock);

    // Test with the fast response
    println!("Requesting fast user...");
    let start = std::time::Instant::now();

    let request = GetUserRequest {
        user_id: "fast-user".to_string(),
    };
    let response = client.get_user(Request::new(request)).await?;

    println!(
        "Fast user response received in {:?}: {}",
        start.elapsed(),
        response.get_ref().name
    );

    // Test with the slow response
    println!("Requesting slow user...");
    let start = std::time::Instant::now();

    let request = GetUserRequest {
        user_id: "slow-user".to_string(),
    };
    let response = client.get_user(Request::new(request)).await?;

    println!(
        "Slow user response received in {:?}: {}",
        start.elapsed(),
        response.get_ref().name
    );

    Ok(())
}