chore: return a future.

@@ -126,14 +126,10 @@ public void testMaintainMultiplexedSession_whenNewSessionCreated_assertThatStale
 // Run second maintenance loop. the first session would now be stale since it has now existed
 // for more than 7 days.
 runMaintenanceLoop(clock, pool, 1);
-// Add a delay is session is swapped asynchronously
-while (multiplexedSessionsRemoved.isEmpty()) {
-Thread.sleep(1);
-}
 SessionFutureWrapper session2 = pool.getMultiplexedSessionWithFallback();
 assertNotEquals(session1.get().getName(), session2.get().getName());
 assertEquals(1, multiplexedSessionsRemoved.size());
-assertTrue(multiplexedSessionsRemoved.contains(session1.get()));
+assertTrue(multiplexedSessionsRemoved.contains(session1.get().get()));
 
 // Advance clock by 8 days
 clock.currentTimeMillis.addAndGet(Duration.ofDays(8).toMillis());
@@ -145,7 +141,7 @@ public void testMaintainMultiplexedSession_whenNewSessionCreated_assertThatStale
 SessionFutureWrapper session3 = pool.getMultiplexedSessionWithFallback();
 assertNotEquals(session2.get().getName(), session3.get().getName());
 assertEquals(2, multiplexedSessionsRemoved.size());
-assertTrue(multiplexedSessionsRemoved.contains(session2.get()));
+assertTrue(multiplexedSessionsRemoved.contains(session2.get().get()));
 }
 
 @Test
@@ -255,7 +251,7 @@ public void testMaintainMultiplexedSession_whenNewSessionCreated_assertThatStale
 // the last attempt.
 runMaintenanceLoop(clock, pool, 1);
 SessionFutureWrapper session3 = pool.getMultiplexedSessionWithFallback();
-assertTrue(multiplexedSessionsRemoved.contains(session1.get()));
+assertTrue(multiplexedSessionsRemoved.contains(session1.get().get()));
 assertNotEquals(session1.get().getName(), session3.get().getName());
 verify(sessionClient, times(2))
 .createMultiplexedSession(any(MultiplexedSessionMaintainerConsumer.class));