From fd39182841daeb1295367316724b7da976959b00 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Wed, 12 Feb 2025 18:02:21 -0800
Subject: [PATCH] fix int tests

---
 crates/prompt_gateway/tests/integration.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/crates/prompt_gateway/tests/integration.rs b/crates/prompt_gateway/tests/integration.rs
index 4af98166..0792a319 100644
--- a/crates/prompt_gateway/tests/integration.rs
+++ b/crates/prompt_gateway/tests/integration.rs
@@ -81,10 +81,11 @@ fn normal_flow(module: &mut Tester, filter_context: i32, http_context: i32) {
                 (":path", "/function_calling"),
                 ("content-type", "application/json"),
                 (":authority", "model_server"),
+                ("x-envoy-upstream-rq-timeout-ms", "30000"),
             ]),
             None,
             None,
-            None,
+            Some(5000),
         )
         .returning(Some(1))
         .expect_log(Some(LogLevel::Trace), None)
@@ -387,10 +388,11 @@ fn prompt_gateway_request_to_llm_gateway() {
                 (":authority", "api_server"),
                 ("x-envoy-max-retries", "3"),
                 (":path", "/weather"),
+                ("x-envoy-upstream-rq-timeout-ms", "30000"),
             ]),
             Some(expected_body),
             None,
-            None,
+            Some(5000),
         )
         .returning(Some(2))
         .expect_metric_increment("active_http_calls", 1)