Rollback unnecessary changes to in_memory_channel_resolver.rs

Author: gabotechs

src/test_utils/in_memory_channel_resolver.rs

@@ -1,7 +1,6 @@
 use crate::{
     ArrowFlightEndpoint, BoxCloneSyncChannel, ChannelResolver, DistributedExt,
-    DistributedSessionBuilder, DistributedSessionBuilderContext,
-    MappedDistributedSessionBuilderExt,
+    DistributedSessionBuilderContext,
 };
 use arrow_flight::flight_service_server::FlightServiceServer;
 use async_trait::async_trait;
@@ -21,17 +20,12 @@ pub struct InMemoryChannelResolver {
 
 impl Default for InMemoryChannelResolver {
     fn default() -> Self {
-        Self::new(|ctx: DistributedSessionBuilderContext| async move {
-            let builder = SessionStateBuilder::new()
-                .with_default_features()
-                .with_runtime_env(ctx.runtime_env.clone());
-            Ok(builder.build())
-        })
+        Self::new()
     }
 }
 
 impl InMemoryChannelResolver {
-    pub fn new(builder: impl DistributedSessionBuilder + Send + Sync + 'static) -> Self {
+    pub fn new() -> Self {
         let (client, server) = tokio::io::duplex(1024 * 1024);
 
         let mut client = Some(client);
@@ -49,10 +43,18 @@ impl InMemoryChannelResolver {
         };
         let this_clone = this.clone();
 
-        let builder =
-            builder.map(move |b| Ok(b.with_distributed_channel_resolver(this.clone()).build()));
-
-        let endpoint = ArrowFlightEndpoint::try_new(builder).unwrap();
+        let endpoint =
+            ArrowFlightEndpoint::try_new(move |ctx: DistributedSessionBuilderContext| {
+                let this = this.clone();
+                async move {
+                    let builder = SessionStateBuilder::new()
+                        .with_default_features()
+                        .with_distributed_channel_resolver(this)
+                        .with_runtime_env(ctx.runtime_env.clone());
+                    Ok(builder.build())
+                }
+            })
+            .unwrap();
 
         tokio::spawn(async move {
             Server::builder()