refactor(grpc-server): move test auth logic to a single function

Author: TanmayArya-1p

crates/grpc_server/src/constants.rs

@@ -12,3 +12,4 @@ pub const ENV_LOGGING: &str = "GRPC_SERVER_LOGGING";
 pub const DEFAULT_PORT: &str = "8080";
 
 pub const SIMILARITY_PROTOBUFF_MAP: [Similarity; 4] = [Euclidean, Manhattan, Hamming, Cosine];
+pub const AUTHORIZATION_HEADER_KEY: &str = "authorization";

crates/grpc_server/src/interceptors.rs

@@ -1,14 +1,16 @@
 use tonic::{Status, service::Interceptor};
 use tracing::{Level, event};
 
+use crate::constants::AUTHORIZATION_HEADER_KEY;
+
 #[derive(Clone)]
 pub struct AuthInterceptor {
     root_password: String,
 }
 
 impl Interceptor for AuthInterceptor {
     fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, Status> {
-        let auth_token = match req.metadata().get("authorization") {
+        let auth_token = match req.metadata().get(AUTHORIZATION_HEADER_KEY) {
             Some(t) => t,
             None => return Err(Status::unauthenticated("Invalid credentials")),
         };

crates/grpc_server/src/tests.rs

@@ -1,4 +1,5 @@
 use crate::config::GRPCServerConfig;
+use crate::constants::AUTHORIZATION_HEADER_KEY;
 use crate::service::vectordb::vector_db_client::VectorDbClient;
 use crate::service::vectordb::{DenseVector, InsertVectorRequest, PointId, SearchRequest};
 use crate::service::{VectorDBService, run_server};
@@ -15,6 +16,15 @@ use tonic::transport::Channel;
 
 // Inspired from https://github.com/hyperium/tonic/discussions/924#discussioncomment-9854088
 
+const TEST_AUTH_BEARER_TOKEN: &str = "123";
+
+fn append_test_auth_header<T>(request: &mut tonic::Request<T>, token: &str) {
+    let auth_value = format!("Bearer {}", token);
+    request
+        .metadata_mut()
+        .insert(AUTHORIZATION_HEADER_KEY, auth_value.parse().unwrap());
+}
+
 async fn start_test_server() -> Result<SocketAddr, Box<dyn std::error::Error>> {
     // using a temporary directory for db datapath
     let temp_dir = tempdir().unwrap();
@@ -28,7 +38,7 @@ async fn start_test_server() -> Result<SocketAddr, Box<dyn std::error::Error>> {
 
     let config = GRPCServerConfig {
         addr: "127.0.0.1:0".parse()?,
-        root_password: "123".to_string(),
+        root_password: TEST_AUTH_BEARER_TOKEN.to_string(),
         logging: false,
         db_config,
     };
@@ -80,9 +90,8 @@ async fn test_grpc_server_start() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let status = client.insert_vector(request).await.is_ok();
     assert!(status);
 }
@@ -101,9 +110,8 @@ async fn test_insert_vector_rpc() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.insert_vector(request).await;
 
     // check if request is successful
@@ -113,9 +121,7 @@ async fn test_insert_vector_rpc() {
     let mut request = tonic::Request::new(PointId {
         id: resp.unwrap().into_inner().id,
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
     let resp = client.get_point(request).await;
 
     // check if request is successful
@@ -130,9 +136,8 @@ async fn test_insert_vector_rpc() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.insert_vector(request).await;
 
     // request must fail
@@ -152,9 +157,8 @@ async fn test_delete_vector_rpc() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.insert_vector(request).await;
 
     // check if request is successful
@@ -163,19 +167,17 @@ async fn test_delete_vector_rpc() {
 
     // delete the vector
     let mut request = tonic::Request::new(PointId { id: point.id });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.delete_point(request).await;
 
     // check if request is successful
     assert!(resp.is_ok());
 
     // verify that the vector is deleted
     let mut request = tonic::Request::new(PointId { id: point.id });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.get_point(request).await;
 
     // request must fail since the vector is deleted
@@ -195,9 +197,8 @@ async fn test_search_vector_rpc() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.insert_vector(request).await;
 
     // check if request is successful
@@ -214,9 +215,8 @@ async fn test_search_vector_rpc() {
         similarity: 0, // euclidean distance
         limit: 1,
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 123".parse().unwrap());
+    append_test_auth_header(&mut request, TEST_AUTH_BEARER_TOKEN);
+
     let resp = client.search_points(request).await;
 
     // check if request is successful
@@ -243,9 +243,8 @@ async fn test_unauthorized_rpc() {
         }),
         payload: Some(Struct::default()),
     });
-    request
-        .metadata_mut()
-        .insert("authorization", "Bearer 43121".parse().unwrap()); // wrong auth token
+
+    append_test_auth_header(&mut request, "43121");
     let resp = client.insert_vector(request).await;
 
     // request must fail