From dbabfc224c05a633c88d1dac69ea2b83d9981a66 Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@pivotal.io>
Date: Thu, 15 Jun 2017 11:29:31 +0200
Subject: [PATCH] Add support for Lettuce in RedisTestServer

Previously, RedisTestServer only supported Jedis and would blow up if
only Lettuce was on the classpath. This commit defensively checks which
driver is available and chose the appropriate one, defaulting to Jedis.

Closes gh-9524
---
 spring-boot-test-autoconfigure/pom.xml        | 10 +--
 .../testsupport/rule/RedisTestServer.java     | 62 ++++++++++++++++---
 2 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/spring-boot-test-autoconfigure/pom.xml b/spring-boot-test-autoconfigure/pom.xml
index b56faa87b9..51d395e1e6 100644
--- a/spring-boot-test-autoconfigure/pom.xml
+++ b/spring-boot-test-autoconfigure/pom.xml
@@ -173,6 +173,11 @@
 			<artifactId>unboundid-ldapsdk</artifactId>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>io.lettuce</groupId>
+			<artifactId>lettuce-core</artifactId>
+			<scope>test</scope>
+		</dependency>
 		<dependency>
 			<groupId>io.projectreactor</groupId>
 			<artifactId>reactor-core</artifactId>
@@ -243,11 +248,6 @@
 			<artifactId>tomcat-embed-el</artifactId>
 			<scope>test</scope>
 		</dependency>
-		<dependency>
-			<groupId>redis.clients</groupId>
-			<artifactId>jedis</artifactId>
-			<scope>test</scope>
-		</dependency>
 		<dependency>
 			<groupId>de.flapdoodle.embed</groupId>
 			<artifactId>de.flapdoodle.embed.mongo</artifactId>
diff --git a/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/rule/RedisTestServer.java b/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/rule/RedisTestServer.java
index a077fa50b1..e6f8a98d84 100644
--- a/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/rule/RedisTestServer.java
+++ b/spring-boot-tools/spring-boot-test-support/src/main/java/org/springframework/boot/testsupport/rule/RedisTestServer.java
@@ -16,6 +16,8 @@
 
 package org.springframework.boot.testsupport.rule;
 
+import java.time.Duration;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.junit.Assume;
@@ -23,8 +25,13 @@ import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
+import org.springframework.beans.factory.DisposableBean;
 import org.springframework.data.redis.connection.RedisConnectionFactory;
+import org.springframework.data.redis.connection.RedisStandaloneConfiguration;
 import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;
+import org.springframework.data.redis.connection.lettuce.LettuceClientConfiguration;
+import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
+import org.springframework.util.ClassUtils;
 
 /**
  * {@link TestRule} for working with an optional Redis server.
@@ -38,7 +45,7 @@ public class RedisTestServer implements TestRule {
 
 	private static final Log logger = LogFactory.getLog(RedisTestServer.class);
 
-	private JedisConnectionFactory connectionFactory;
+	private RedisConnectionFactory connectionFactory;
 
 	@Override
 	public Statement apply(final Statement base, Description description) {
@@ -52,14 +59,23 @@ public class RedisTestServer implements TestRule {
 		}
 	}
 
-	private JedisConnectionFactory createConnectionFactory() {
-		JedisConnectionFactory connectionFactory = new JedisConnectionFactory();
-		connectionFactory.afterPropertiesSet();
-		testConnection(connectionFactory);
-		return connectionFactory;
+	private RedisConnectionFactory createConnectionFactory() {
+		ClassLoader classLoader = RedisTestServer.class.getClassLoader();
+		RedisConnectionFactory cf;
+		if (ClassUtils.isPresent("redis.clients.jedis.Jedis", classLoader)) {
+			cf = new JedisConnectionFactoryConfiguration()
+					.createConnectionFactory();
+		}
+		else {
+			cf = new LettuceConnectionFactoryConfiguration()
+					.createConnectionFactory();
+		}
+
+		testConnection(cf);
+		return cf;
 	}
 
-	private void testConnection(JedisConnectionFactory connectionFactory) {
+	private void testConnection(RedisConnectionFactory connectionFactory) {
 		connectionFactory.getConnection().close();
 	}
 
@@ -76,9 +92,9 @@ public class RedisTestServer implements TestRule {
 
 		private final Statement base;
 
-		private final JedisConnectionFactory connectionFactory;
+		private final RedisConnectionFactory connectionFactory;
 
-		RedisStatement(Statement base, JedisConnectionFactory connectionFactory) {
+		RedisStatement(Statement base, RedisConnectionFactory connectionFactory) {
 			this.base = base;
 			this.connectionFactory = connectionFactory;
 		}
@@ -90,7 +106,9 @@ public class RedisTestServer implements TestRule {
 			}
 			finally {
 				try {
-					this.connectionFactory.destroy();
+					if (this.connectionFactory instanceof DisposableBean) {
+						((DisposableBean) this.connectionFactory).destroy();
+					}
 				}
 				catch (Exception ex) {
 					logger.warn("Exception while trying to cleanup redis resource", ex);
@@ -110,4 +128,28 @@ public class RedisTestServer implements TestRule {
 
 	}
 
+	private static class JedisConnectionFactoryConfiguration {
+
+		RedisConnectionFactory createConnectionFactory() {
+			JedisConnectionFactory connectionFactory = new JedisConnectionFactory();
+			connectionFactory.afterPropertiesSet();
+			return connectionFactory;
+		}
+
+	}
+
+	private static class LettuceConnectionFactoryConfiguration {
+
+		RedisConnectionFactory createConnectionFactory() {
+			LettuceClientConfiguration config = LettuceClientConfiguration.builder()
+					.shutdownTimeout(Duration.ofMillis(0))
+					.build();
+			LettuceConnectionFactory connectionFactory = new LettuceConnectionFactory(
+					new RedisStandaloneConfiguration(), config);
+			connectionFactory.afterPropertiesSet();
+			return connectionFactory;
+		}
+
+	}
+
 }