refactor: attempt to create constructor that allows passing in services

1 files changed, 22 insertions, 18 deletions

diff --git a/src/authorization/server.rs b/src/authorization/server.rs
index f11d0465..7c39b51c 100644
--- a/src/authorization/server.rs
+++ b/src/authorization/server.rs
@@ -3,30 +3,34 @@ use super::check_service::CheckService;
 use envoy_types::ext_authz::v3::pb::AuthorizationServer;
 use std::sync::Arc;
 
-pub fn create_router() -> Result<tonic::transport::server::Router, Box<dyn std::error::Error>> {
-    let (_health_reporter, health_service) = tonic_health::server::health_reporter();
-    let authorizer = Arc::new(CedarAuthorizer::default());
-    let check_service = CheckService::new(authorizer);
-    let server = tonic::transport::Server::builder()
-        .add_service(AuthorizationServer::new(check_service))
-        .add_service(health_service)
-        .add_service(
-            tonic_reflection::server::Builder::configure()
-                .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
-                .build_v1()
-                .unwrap(),
-        );
-    Ok(server)
-}
-
 pub struct Server {
     router: tonic::transport::server::Router,
 }
 
 impl Server {
     pub fn new() -> Result<Server, Box<dyn std::error::Error>> {
-        let router = create_router()?;
-        Ok(Server { router: router })
+        Ok(Self::new_with(|mut builder| {
+            let (_health_reporter, health_service) = tonic_health::server::health_reporter();
+            let authorizer = Arc::new(CedarAuthorizer::default());
+            let check_service = CheckService::new(authorizer);
+            builder
+                .add_service(AuthorizationServer::new(check_service))
+                .add_service(health_service)
+                .add_service(
+                    tonic_reflection::server::Builder::configure()
+                        .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
+                        .build_v1()
+                        .unwrap(),
+                )
+        }))
+    }
+
+    pub fn new_with<F>(f: F) -> Server
+    where
+        F: FnOnce(tonic::transport::Server) -> tonic::transport::server::Router,
+    {
+        let router = f(tonic::transport::Server::builder());
+        Server { router }
     }
 
     pub async fn serve(self, addr: std::net::SocketAddr) -> Result<(), tonic::transport::Error> {