Add helpers to Jedis pool (#4366)

* Add helpers to Jedis pool

* Test for helpers to Jedis pool

* amedn method and test

* Clean up + unit test

   - Add unit test to verify connection is returned to the pool

---------

Co-authored-by: ggivo <[email protected]>

Author: oscar-besga-panel

pom.xml

@@ -546,6 +546,7 @@
 						<include>**/*CommandFlags*.java</include>
 						<include>**/*CompareCondition*.java</include>
 						<include>**/*MSetExParams*.java</include>
+						<include>**/*UnitTest.java</include>
 					</includes>
 				</configuration>
 				<executions>

src/main/java/redis/clients/jedis/JedisPool.java

@@ -1,6 +1,8 @@
 package redis.clients.jedis;
 
 import java.net.URI;
+import java.util.function.Consumer;
+import java.util.function.Function;
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLSocketFactory;
@@ -392,4 +394,17 @@ public void returnResource(final Jedis resource) {
       }
     }
   }
+
+  public void withResource(Consumer<Jedis> consumer) {
+    try (Jedis jedis = this.getResource()) {
+      consumer.accept(jedis);
+    }
+  }
+
+  public <K> K withResourceGet(Function<Jedis, K> function) {
+    try (Jedis jedis = this.getResource()) {
+      return function.apply(jedis);
+    }
+  }
+
 }

src/test/java/redis/clients/jedis/JedisPoolTest.java

@@ -9,15 +9,19 @@
 import org.apache.commons.pool2.impl.DefaultPooledObject;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
 import redis.clients.jedis.exceptions.InvalidURIException;
 import redis.clients.jedis.exceptions.JedisAccessControlException;
 import redis.clients.jedis.exceptions.JedisConnectionException;
 import redis.clients.jedis.exceptions.JedisException;
 
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.awaitility.Awaitility.await;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -33,6 +37,15 @@ public class JedisPoolTest {
 
   private static final EndpointConfig endpointStandalone1 = HostAndPorts.getRedisEndpoint("standalone1");
 
+  private String testKey;
+  private String testValue;
+
+  @BeforeEach
+  public void setUpTestKey(TestInfo testInfo) {
+    testKey = testInfo.getDisplayName() + "-key";
+    testValue = testInfo.getDisplayName() + "-value";
+  }
+
   @Test
   public void checkConnections() {
     JedisPool pool = new JedisPool(new JedisPoolConfig(), endpointStandalone0.getHost(), endpointStandalone0.getPort(), 2000);
@@ -462,4 +475,64 @@ public void testResetValidCredentials() {
       }
     }
   }
+
+  @Test
+  public void testWithResource() {
+    try (JedisPool pool = new JedisPool(new JedisPoolConfig(), endpointStandalone0.getHostAndPort(),
+        endpointStandalone0.getClientConfigBuilder().build())) {
+
+      pool.withResource(jedis -> {
+        jedis.set(testKey, testValue);
+      });
+
+      pool.withResource(jedis -> {
+        assertEquals(testValue, jedis.get(testKey));
+      });
+    }
+  }
+
+  @Test
+  public void testWithResourceReturnsConnectionToPool() {
+    try (JedisPool pool = new JedisPool(new JedisPoolConfig(), endpointStandalone0.getHostAndPort(),
+        endpointStandalone0.getClientConfigBuilder().build())) {
+
+      pool.withResource(jedis -> {
+        assertThat(pool.getNumActive(), equalTo(1));
+        jedis.set("foo", "bar");
+      });
+
+      assertThat(pool.getNumActive(), equalTo(0));
+    }
+  }
+
+  @Test
+  public void testWithResourceGet() {
+    try (JedisPool pool = new JedisPool(new JedisPoolConfig(), endpointStandalone0.getHostAndPort(),
+        endpointStandalone0.getClientConfigBuilder().build())) {
+
+      String result = pool.withResourceGet(jedis -> {
+        jedis.set(testKey, testValue);
+        return jedis.get(testKey);
+      });
+
+      assertEquals(testValue, result);
+    }
+  }
+
+  @Test
+  public void testWithResourceGetReturnsConnectionToPool() {
+    try (JedisPool pool = new JedisPool(new JedisPoolConfig(), endpointStandalone0.getHostAndPort(),
+        endpointStandalone0.getClientConfigBuilder().build())) {
+
+      String result = pool.withResourceGet(jedis -> {
+        assertThat(pool.getNumActive(), equalTo(1));
+        jedis.set("foo", "bar");
+        return jedis.get("foo");
+      });
+
+      assertThat(result, equalTo("bar"));
+      assertThat(pool.getNumActive(), equalTo(0));
+    }
+  }
+
 }

src/test/java/redis/clients/jedis/JedisPoolUnitTest.java

@@ -0,0 +1,60 @@
+package redis.clients.jedis;
+
+import org.apache.commons.pool2.PooledObjectFactory;
+import org.apache.commons.pool2.impl.DefaultPooledObject;
+import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
+
+public class JedisPoolUnitTest {
+
+  private JedisPool pool;
+  private Jedis mockJedis;
+
+  @BeforeEach
+  public void setUp() throws Exception {
+    mockJedis = mock(Jedis.class);
+    PooledObjectFactory<Jedis> mockFactory = mock(PooledObjectFactory.class);
+
+    when(mockFactory.makeObject()).thenReturn(new DefaultPooledObject<>(mockJedis));
+
+    GenericObjectPoolConfig<Jedis> config = new GenericObjectPoolConfig<>();
+    config.setMaxTotal(1);
+    pool = spy(new JedisPool(config, mockFactory));
+
+  }
+
+  @AfterEach
+  public void tearDown() {
+    if (pool != null && !pool.isClosed()) {
+      pool.close();
+    }
+  }
+
+  @Test
+  public void testWithResourceClosesConnection() {
+    pool.withResource(jedis -> assertEquals(mockJedis, jedis));
+
+    verify(mockJedis, times(1)).close();
+  }
+
+  @Test
+  public void testWithResourceGetClosesConnection() {
+    String result = pool.withResourceGet(jedis -> "test-result");
+
+    verify(mockJedis, times(1)).close();
+  }
+
+  @Test
+  public void testWithResourceGetReturnsResult() {
+    when(mockJedis.get(eq("test-key"))).thenReturn("test-result");
+    String result = pool.withResourceGet(jedis -> jedis.get("test-key"));
+
+    verify(mockJedis, times(1)).close();
+  }
+
+}