mcp-server/tests/common/mock_model_server.rs
48 lines
use std::sync::atomic::AtomicUsize;use std::sync::atomic::Ordering;use wiremock::Mock;use wiremock::MockServer;use wiremock::Respond;use wiremock::ResponseTemplate;use wiremock::matchers::method;use wiremock::matchers::path;/// Create a mock server that will provide the responses, in order, for/// requests to the `/v1/responses` endpoint.pub async fn create_mock_responses_server(responses: Vec<String>) -> MockServer { let server = MockServer::start().await; let num_calls = responses.len(); let seq_responder = SeqResponder { num_calls: AtomicUsize::new(0), responses, }; Mock::given(method("POST")) .and(path("/v1/responses")) .respond_with(seq_responder) .expect(num_calls as u64) .mount(&server) .await; server}struct SeqResponder { num_calls: AtomicUsize, responses: Vec<String>,}impl Respond for SeqResponder { fn respond(&self, _: &wiremock::Request) -> ResponseTemplate { let call_num = self.num_calls.fetch_add(1, Ordering::SeqCst); let response = self .responses .get(call_num) .expect("mock model response should exist"); ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") .set_body_raw(response.clone(), "text/event-stream") }}