From 270ff1f5decf3891a87febf52c965f00dfa15e27 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil@katanemo.com>
Date: Thu, 7 Nov 2024 14:35:43 -0800
Subject: [PATCH] fix tests

---
 crates/prompt_gateway/src/stream_context.rs | 6 ------
 crates/prompt_gateway/tests/integration.rs  | 2 ++
 2 files changed, 2 insertions(+), 6 deletions(-)

diff --git a/crates/prompt_gateway/src/stream_context.rs b/crates/prompt_gateway/src/stream_context.rs
index f92c4ad5..dac50ad5 100644
--- a/crates/prompt_gateway/src/stream_context.rs
+++ b/crates/prompt_gateway/src/stream_context.rs
@@ -170,7 +170,6 @@ impl StreamContext {
             ("content-type", "application/json"),
             ("x-envoy-max-retries", "3"),
             ("x-envoy-upstream-rq-timeout-ms", "60000"),
-            ("x-envoy-force-trace", "true"),
         ];
 
         if self.request_id.is_some() {
@@ -300,7 +299,6 @@ impl StreamContext {
             ("content-type", "application/json"),
             ("x-envoy-max-retries", "3"),
             ("x-envoy-upstream-rq-timeout-ms", "60000"),
-            ("x-envoy-force-trace", "true"),
         ];
 
         if self.request_id.is_some() {
@@ -511,7 +509,6 @@ impl StreamContext {
                     ("content-type", "application/json"),
                     ("x-envoy-max-retries", "3"),
                     ("x-envoy-upstream-rq-timeout-ms", timeout_str.as_str()),
-                    ("x-envoy-force-trace", "true"),
                 ];
 
                 if self.request_id.is_some() {
@@ -663,7 +660,6 @@ impl StreamContext {
             ("content-type", "application/json"),
             ("x-envoy-max-retries", "3"),
             ("x-envoy-upstream-rq-timeout-ms", timeout_str.as_str()),
-            ("x-envoy-force-trace", "true"),
         ];
 
         if self.request_id.is_some() {
@@ -844,7 +840,6 @@ impl StreamContext {
             ("content-type", "application/json"),
             ("x-envoy-max-retries", "3"),
             ("x-envoy-upstream-rq-timeout-ms", "60000"),
-            ("x-envoy-force-trace", "true"),
         ];
 
         if self.request_id.is_some() {
@@ -900,7 +895,6 @@ impl StreamContext {
             (":authority", endpoint.name.as_str()),
             ("content-type", "application/json"),
             ("x-envoy-max-retries", "3"),
-            ("x-envoy-force-trace", "true"),
         ];
 
         if self.request_id.is_some() {
diff --git a/crates/prompt_gateway/tests/integration.rs b/crates/prompt_gateway/tests/integration.rs
index 46f2dfd8..821c3543 100644
--- a/crates/prompt_gateway/tests/integration.rs
+++ b/crates/prompt_gateway/tests/integration.rs
@@ -36,6 +36,8 @@ fn request_headers_expectations(module: &mut Tester, http_context: i32) {
         .expect_log(Some(LogLevel::Trace), 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)
         .execute_and_expect(ReturnType::Action(Action::Continue))
         .unwrap();
 }