2 files changed, 12 insertions, 17 deletions

diff --git a/tests/common/factory_bot.rs b/tests/common/factory_bot.rs
index f25633f4..2389c858 100644
--- a/tests/common/factory_bot.rs
+++ b/tests/common/factory_bot.rs
@@ -26,14 +26,6 @@ pub fn create_request(f: impl std::ops::FnOnce(&mut HttpRequest)) -> CheckReques
     })
 }
 
-pub fn create_test_request_with_headers(
-    headers: HashMap<String, String>,
-) -> tonic::Request<CheckRequest> {
-    tonic::Request::new(create_request(|item: &mut HttpRequest| {
-        item.headers = headers;
-    }))
-}
-
 pub fn build_headers(headers: Vec<(String, String)>) -> HashMap<String, String> {
     return build_with(|item: &mut HashMap<String, String>| {
         for (key, value) in headers {
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index c25d1920..56321acb 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -1,5 +1,7 @@
 use authzd::{CedarAuthorizer, CheckService};
+use common::*;
 use envoy_types::ext_authz::v3::pb::Authorization;
+use envoy_types::pb::envoy::service::auth::v3::attribute_context::HttpRequest;
 use std::sync::Arc;
 
 mod authorization;
@@ -9,12 +11,12 @@ mod common;
 async fn test_success_response() {
     let authorizer = Arc::new(CedarAuthorizer::new());
     let server = CheckService::new(authorizer);
-    let request = common::factory_bot::create_test_request_with_headers(
-        common::factory_bot::build_headers(vec![(
+    let request = tonic::Request::new(factory_bot::create_request(|item: &mut HttpRequest| {
+        item.headers = factory_bot::build_headers(vec![(
             "authorization".to_string(),
             "Bearer valid-token".to_string(),
-        )]),
-    );
+        )])
+    }));
 
     let response = server.check(request).await;
     assert!(response.is_ok());
@@ -39,11 +41,12 @@ async fn test_multiple() {
     ];
 
     for (auth_value, should_succeed) in test_cases {
-        let headers = common::factory_bot::build_headers(vec![(
-            "authorization".to_string(),
-            auth_value.to_string(),
-        )]);
-        let request = common::factory_bot::create_test_request_with_headers(headers);
+        let request = tonic::Request::new(factory_bot::create_request(|item: &mut HttpRequest| {
+            item.headers = factory_bot::build_headers(vec![(
+                "authorization".to_string(),
+                auth_value.to_string(),
+            )]);
+        }));
 
         let response = server.check(request).await;
         assert!(response.is_ok());