From c343d72e6d7baece9a6a1127bad3c2e6e85739c7 Mon Sep 17 00:00:00 2001 From: Adil Hafeez Date: Thu, 30 Jan 2025 19:21:28 -0800 Subject: [PATCH] fix rust tests --- crates/llm_gateway/tests/integration.rs | 46 ++++++++++----------- crates/prompt_gateway/src/context.rs | 5 +-- crates/prompt_gateway/src/stream_context.rs | 5 ++- crates/prompt_gateway/tests/integration.rs | 28 +++++++------ 4 files changed, 41 insertions(+), 43 deletions(-) diff --git a/crates/llm_gateway/tests/integration.rs b/crates/llm_gateway/tests/integration.rs index 73c431bf..0b28a175 100644 --- a/crates/llm_gateway/tests/integration.rs +++ b/crates/llm_gateway/tests/integration.rs @@ -18,23 +18,12 @@ fn wasm_module() -> String { fn request_headers_expectations(module: &mut Tester, http_context: i32) { module .call_proxy_on_request_headers(http_context, 0, false) - .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path")) - .returning(Some("/test_path")) - .expect_log(Some(LogLevel::Debug), Some("request_path: /test_path")) .expect_get_header_map_value( Some(MapType::HttpRequestHeaders), Some("x-arch-llm-provider-hint"), ) - .returning(Some("default")) - .expect_get_header_map_value( - Some(MapType::HttpRequestHeaders), - Some("x-arch-llm-provider-hint"), - ) - .returning(Some("default")) - .expect_log( - Some(LogLevel::Debug), - Some("llm provider hint: Some(\"default\"), selected llm: open-ai-gpt-4"), - ) + .returning(None) + .expect_log(Some(LogLevel::Debug), Some("request received: llm provider hint: Some(\"default\"), selected llm: open-ai-gpt-4")) .expect_add_header_map_value( Some(MapType::HttpRequestHeaders), Some("x-arch-llm-provider"), @@ -45,7 +34,11 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { Some("Authorization"), Some("Bearer secret_key"), ) - .expect_remove_header_map_value(Some(MapType::HttpRequestHeaders), Some("content-length")) + .expect_get_header_map_value( + Some(MapType::HttpRequestHeaders), + Some("x-arch-llm-provider-hint"), + ) + .returning(Some("default")) .expect_get_header_map_value( Some(MapType::HttpRequestHeaders), Some("x-arch-ratelimit-selector"), @@ -57,11 +50,10 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { .returning(None) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some(":path")) .returning(Some("/v1/chat/completions")) - // .expect_log(Some(LogLevel::Debug), None) .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("x-request-id")) .returning(None) - // .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("traceparent")) - // .returning(None) + .expect_get_header_map_value(Some(MapType::HttpRequestHeaders), Some("traceparent")) + .returning(None) .execute_and_expect(ReturnType::Action(Action::Continue)) .unwrap(); } @@ -69,7 +61,7 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { fn normal_flow(module: &mut Tester, filter_context: i32, http_context: i32) { module .call_proxy_on_context_create(http_context, filter_context) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -194,6 +186,10 @@ fn llm_gateway_successful_request_to_open_ai_chat_completions() { module .call_proxy_on_context_create(http_context, filter_context) + .expect_log( + Some(LogLevel::Trace), + Some("||| create_http_context called with context_id: 2 |||"), + ) .execute_and_expect(ReturnType::None) .unwrap(); @@ -224,9 +220,9 @@ fn llm_gateway_successful_request_to_open_ai_chat_completions() { .expect_get_buffer_bytes(Some(BufferType::HttpRequestBody)) .returning(Some(chat_completions_request_body)) .expect_log(Some(LogLevel::Trace), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_metric_record("input_sequence_length", 21) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None) @@ -257,7 +253,7 @@ fn llm_gateway_bad_request_to_open_ai_chat_completions() { module .call_proxy_on_context_create(http_context, filter_context) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -345,9 +341,9 @@ fn llm_gateway_request_ratelimited() { .returning(Some(chat_completions_request_body)) // The actual call is not important in this test, we just need to grab the token_id .expect_log(Some(LogLevel::Trace), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_metric_record("input_sequence_length", 107) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) @@ -411,9 +407,9 @@ fn llm_gateway_request_not_ratelimited() { .returning(Some(chat_completions_request_body)) // The actual call is not important in this test, we just need to grab the token_id .expect_log(Some(LogLevel::Trace), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_metric_record("input_sequence_length", 29) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None) diff --git a/crates/prompt_gateway/src/context.rs b/crates/prompt_gateway/src/context.rs index 5e8fcf0a..89725e0d 100644 --- a/crates/prompt_gateway/src/context.rs +++ b/crates/prompt_gateway/src/context.rs @@ -54,10 +54,7 @@ impl Context for StreamContext { } } else { // :status header not found - return self.send_server_error( - ServerError::LogicError("missing :status header".to_string()), - Some(StatusCode::BAD_GATEWAY), - ); + warn!("missing :status header"); } #[cfg_attr(any(), rustfmt::skip)] diff --git a/crates/prompt_gateway/src/stream_context.rs b/crates/prompt_gateway/src/stream_context.rs index c765cf16..3e3fee5b 100644 --- a/crates/prompt_gateway/src/stream_context.rs +++ b/crates/prompt_gateway/src/stream_context.rs @@ -363,7 +363,10 @@ impl StreamContext { let http_status = self .get_http_call_response_header(":status") .unwrap_or(StatusCode::OK.as_str().to_string()); - debug!("developer api call response received: status code: {}", http_status); + debug!( + "developer api call response received: status code: {}", + http_status + ); if http_status != StatusCode::OK.as_str() { warn!( "api server responded with non 2xx status code: {}", diff --git a/crates/prompt_gateway/tests/integration.rs b/crates/prompt_gateway/tests/integration.rs index 1a6ed0e6..9fcaf74d 100644 --- a/crates/prompt_gateway/tests/integration.rs +++ b/crates/prompt_gateway/tests/integration.rs @@ -41,7 +41,7 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) { fn normal_flow(module: &mut Tester, filter_context: i32, http_context: i32) { module .call_proxy_on_context_create(http_context, filter_context) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -87,8 +87,9 @@ fn normal_flow(module: &mut Tester, filter_context: i32, http_context: i32) { None, ) .returning(Some(1)) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Trace), None) .expect_metric_increment("active_http_calls", 1) .execute_and_expect(ReturnType::Action(Action::Pause)) @@ -203,7 +204,7 @@ fn prompt_gateway_successful_request_to_open_ai_chat_completions() { module .call_proxy_on_context_create(http_context, filter_context) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -234,8 +235,9 @@ fn prompt_gateway_successful_request_to_open_ai_chat_completions() { .expect_get_buffer_bytes(Some(BufferType::HttpRequestBody)) .returning(Some(chat_completions_request_body)) .expect_log(Some(LogLevel::Trace), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Trace), None) .expect_http_call(Some("arch_internal"), None, None, None, None) .returning(Some(4)) @@ -267,7 +269,7 @@ fn prompt_gateway_bad_request_to_open_ai_chat_completions() { module .call_proxy_on_context_create(http_context, filter_context) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::None) .unwrap(); @@ -302,7 +304,7 @@ fn prompt_gateway_bad_request_to_open_ai_chat_completions() { None, None, ) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::Action(Action::Pause)) .unwrap(); } @@ -369,10 +371,11 @@ fn prompt_gateway_request_to_llm_gateway() { .expect_metric_increment("active_http_calls", -1) .expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody)) .returning(Some(&arch_fc_resp_str)) + .expect_log(Some(LogLevel::Warn), None) .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Trace), None) .expect_http_call( @@ -401,13 +404,12 @@ fn prompt_gateway_request_to_llm_gateway() { .expect_get_buffer_bytes(Some(BufferType::HttpCallResponseBody)) .returning(Some(&body_text)) .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) + .expect_log(Some(LogLevel::Debug), None) + .expect_log(Some(LogLevel::Trace), None) + .expect_log(Some(LogLevel::Trace), None) .expect_get_header_map_value(Some(MapType::HttpCallResponseHeaders), Some(":status")) .returning(Some("200")) - .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) - .expect_log(Some(LogLevel::Debug), None) .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None) .execute_and_expect(ReturnType::None) .unwrap();