From 29a5c2fdf0ce6b8515fd958a6dc7d739d1d89b5a Mon Sep 17 00:00:00 2001 From: Adil Hafeez Date: Mon, 28 Oct 2024 14:04:35 -0700 Subject: [PATCH] fix int tests --- crates/llm_gateway/tests/integration.rs | 24 ++++++++-------------- crates/prompt_gateway/src/http_context.rs | 6 +++--- crates/prompt_gateway/tests/integration.rs | 8 ++++---- 3 files changed, 16 insertions(+), 22 deletions(-) diff --git a/crates/llm_gateway/tests/integration.rs b/crates/llm_gateway/tests/integration.rs index cc17c738..5a5ed4a8 100644 --- a/crates/llm_gateway/tests/integration.rs +++ b/crates/llm_gateway/tests/integration.rs @@ -149,14 +149,14 @@ ratelimits: key: selector-key value: selector-value limit: - tokens: 50 + tokens: 100 unit: minute "# } #[test] #[serial] -fn successful_request_to_open_ai_chat_completions() { +fn llm_gateway_successful_request_to_open_ai_chat_completions() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -217,7 +217,7 @@ fn successful_request_to_open_ai_chat_completions() { #[test] #[serial] -fn bad_request_to_open_ai_chat_completions() { +fn llm_gateway_bad_request_to_open_ai_chat_completions() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -279,7 +279,7 @@ fn bad_request_to_open_ai_chat_completions() { #[test] #[serial] -fn request_ratelimited() { +fn llm_gateway_request_ratelimited() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -306,11 +306,11 @@ fn request_ratelimited() { \"messages\": [\ {\ \"role\": \"system\",\ - \"content\": \"You are a poetic assistant, skilled in explaining complex programming concepts with creative flair.\"\ + \"content\": \"You are a helpful poetic assistant!, skilled in explaining complex programming concepts with creative flair. Be sure to be concise and to the point.\"\ },\ {\ \"role\": \"user\",\ - \"content\": \"Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. \"\ + \"content\": \"Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. Compose a poem that explains the concept of recursion in programming. And also summarize it how a 4th graded would understand it.\"\ }\ ],\ \"model\": \"gpt-4\"\ @@ -328,6 +328,7 @@ fn request_ratelimited() { .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_metric_increment("active_http_calls", 1) .expect_send_local_response( Some(StatusCode::TOO_MANY_REQUESTS.as_u16().into()), @@ -342,7 +343,7 @@ fn request_ratelimited() { #[test] #[serial] -fn request_not_ratelimited() { +fn llm_gateway_request_not_ratelimited() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -391,14 +392,7 @@ fn request_not_ratelimited() { .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Debug), None) - // .expect_metric_increment("active_http_calls", 1) - .expect_send_local_response( - Some(StatusCode::TOO_MANY_REQUESTS.as_u16().into()), - None, - None, - None, - ) - .expect_metric_increment("ratelimited_rq", 1) + .expect_set_buffer_bytes(Some(BufferType::HttpRequestBody), None) .execute_and_expect(ReturnType::Action(Action::Continue)) .unwrap(); } diff --git a/crates/prompt_gateway/src/http_context.rs b/crates/prompt_gateway/src/http_context.rs index 596a6a4e..3b014009 100644 --- a/crates/prompt_gateway/src/http_context.rs +++ b/crates/prompt_gateway/src/http_context.rs @@ -33,7 +33,8 @@ impl HttpContext for StreamContext { // manipulate the body in benign ways e.g., compression. self.set_http_request_header("content-length", None); - if self.get_http_request_header(":path").unwrap_or_default() == HEALTHZ_PATH { + let request_path = self.get_http_request_header(":path").unwrap_or_default(); + if request_path == HEALTHZ_PATH { if self.embeddings_store.is_none() { self.send_http_response(503, vec![], None); } else { @@ -42,8 +43,7 @@ impl HttpContext for StreamContext { return Action::Continue; } - self.is_chat_completions_request = - self.get_http_request_header(":path").unwrap_or_default() == CHAT_COMPLETIONS_PATH; + self.is_chat_completions_request = request_path == CHAT_COMPLETIONS_PATH; trace!( "on_http_request_headers S[{}] req_headers={:?}", diff --git a/crates/prompt_gateway/tests/integration.rs b/crates/prompt_gateway/tests/integration.rs index 27eac427..1bf581c5 100644 --- a/crates/prompt_gateway/tests/integration.rs +++ b/crates/prompt_gateway/tests/integration.rs @@ -375,7 +375,7 @@ ratelimits: #[test] #[serial] -fn successful_request_to_open_ai_chat_completions() { +fn prompt_gateway_successful_request_to_open_ai_chat_completions() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -438,7 +438,7 @@ fn successful_request_to_open_ai_chat_completions() { #[test] #[serial] -fn bad_request_to_open_ai_chat_completions() { +fn prompt_gateway_bad_request_to_open_ai_chat_completions() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -501,7 +501,7 @@ fn bad_request_to_open_ai_chat_completions() { #[test] #[serial] -fn request_to_llm_gateway() { +fn prompt_gateway_request_to_llm_gateway() { let args = tester::MockSettings { wasm_path: wasm_module(), quiet: false, @@ -669,8 +669,8 @@ fn request_to_llm_gateway() { .expect_get_buffer_bytes(Some(BufferType::HttpResponseBody)) .returning(Some(chat_completion_response_str.as_str())) .expect_log(Some(LogLevel::Trace), None) + .expect_log(Some(LogLevel::Debug), None) .expect_set_buffer_bytes(Some(BufferType::HttpResponseBody), None) - .expect_log(Some(LogLevel::Trace), None) .expect_log(Some(LogLevel::Debug), None) .expect_log(Some(LogLevel::Trace), None) .execute_and_expect(ReturnType::Action(Action::Continue))