Polish

See gh-19802

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/data/cassandra/CassandraDataAutoConfigurationIntegrationTests.java

@@ -16,12 +16,10 @@
 
 package org.springframework.boot.autoconfigure.data.cassandra;
 
-import java.net.InetSocketAddress;
 import java.time.Duration;
 
 import com.datastax.oss.driver.api.core.CqlSession;
-import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
-import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
+import com.datastax.oss.driver.api.core.CqlSessionBuilder;
 import org.junit.jupiter.api.Test;
 import org.testcontainers.containers.CassandraContainer;
 import org.testcontainers.junit.jupiter.Container;
@@ -33,6 +31,8 @@ import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.cassandra.CassandraAutoConfiguration;
 import org.springframework.boot.autoconfigure.data.cassandra.city.City;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
 import org.springframework.data.cassandra.config.SchemaAction;
 import org.springframework.data.cassandra.config.SessionFactoryFactoryBean;
 
@@ -68,25 +68,24 @@ class CassandraDataAutoConfigurationIntegrationTests {
 
 	@Test
 	void hasRecreateSchemaActionSet() {
-		createTestKeyspaceIfNotExists();
-		this.contextRunner
-				.withPropertyValues("spring.data.cassandra.schemaAction=recreate_drop_unused",
-						"spring.data.cassandra.keyspaceName=boot_test")
+		this.contextRunner.withUserConfiguration(KeyspaceTestConfiguration.class)
+				.withPropertyValues("spring.data.cassandra.schemaAction=recreate_drop_unused")
 				.run((context) -> assertThat(context.getBean(SessionFactoryFactoryBean.class))
 						.hasFieldOrPropertyWithValue("schemaAction", SchemaAction.RECREATE_DROP_UNUSED));
 	}
 
-	private void createTestKeyspaceIfNotExists() {
-		try (CqlSession session = CqlSession.builder()
-				.withConfigLoader(DriverConfigLoader.programmaticBuilder()
-						.withDuration(DefaultDriverOption.CONNECTION_INIT_QUERY_TIMEOUT, Duration.ofSeconds(10))
-						.withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(20)).build())
-				.addContactPoint(
-						new InetSocketAddress(cassandra.getContainerIpAddress(), cassandra.getFirstMappedPort()))
-				.withLocalDatacenter("datacenter1").build()) {
-			session.execute("CREATE KEYSPACE IF NOT EXISTS boot_test"
-					+ "  WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };");
+	@Configuration(proxyBeanMethods = false)
+	static class KeyspaceTestConfiguration {
+
+		@Bean
+		CqlSession cqlSession(CqlSessionBuilder cqlSessionBuilder) {
+			try (CqlSession session = cqlSessionBuilder.build()) {
+				session.execute("CREATE KEYSPACE IF NOT EXISTS boot_test"
+						+ "  WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 };");
+			}
+			return cqlSessionBuilder.withKeyspace("boot_test").build();
 		}
+
 	}
 
 }