From 320dd0e24e5c07f233ebdf653e0669a05542ad46 Mon Sep 17 00:00:00 2001
From: anessi <anessi@users.noreply.github.com>
Date: Tue, 5 Sep 2023 10:54:55 +0200
Subject: [PATCH] Add virtual host support for Rabbit Stream

Add a new property 'spring.rabbitmq.stream.virtual-host' which can be
used to set a custom virtual host for streams.

See gh-37189
---
 .../autoconfigure/amqp/RabbitProperties.java   | 14 ++++++++++++++
 .../amqp/RabbitStreamConfiguration.java        |  4 ++++
 .../amqp/RabbitStreamConfigurationTests.java   | 18 ++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitProperties.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitProperties.java
index 4e95d7346b..a9ae5efc41 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitProperties.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitProperties.java
@@ -1206,6 +1206,12 @@ public class RabbitProperties {
 		 */
 		private int port = DEFAULT_STREAM_PORT;
 
+		/**
+		 * Virtual host of a RabbitMQ instance with the Stream plugin enabled. When not
+		 * set, spring.rabbitmq.virtual-host is used.
+		 */
+		private String virtualHost;
+
 		/**
 		 * Login user to authenticate to the broker. When not set,
 		 * spring.rabbitmq.username is used.
@@ -1239,6 +1245,14 @@ public class RabbitProperties {
 			this.port = port;
 		}
 
+		public String getVirtualHost() {
+			return this.virtualHost;
+		}
+
+		public void setVirtualHost(String virtualHost) {
+			this.virtualHost = virtualHost;
+		}
+
 		public String getUsername() {
 			return this.username;
 		}
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfiguration.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfiguration.java
index 6547cfdc4e..dfacce8455 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfiguration.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfiguration.java
@@ -102,6 +102,10 @@ class RabbitStreamConfiguration {
 		PropertyMapper map = PropertyMapper.get();
 		map.from(stream.getHost()).to(builder::host);
 		map.from(stream.getPort()).to(builder::port);
+		map.from(stream.getVirtualHost())
+			.as(withFallback(properties::getVirtualHost))
+			.whenNonNull()
+			.to(builder::virtualHost);
 		map.from(stream.getUsername()).as(withFallback(properties::getUsername)).whenNonNull().to(builder::username);
 		map.from(stream.getPassword()).as(withFallback(properties::getPassword)).whenNonNull().to(builder::password);
 		return builder;
diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfigurationTests.java
index eec28db57a..95549628d1 100644
--- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfigurationTests.java
+++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/amqp/RabbitStreamConfigurationTests.java
@@ -143,6 +143,24 @@ class RabbitStreamConfigurationTests {
 		then(builder).should().host("stream.rabbit.example.com");
 	}
 
+	@Test
+	void whenStreamVirtualHostIsSetThenEnvironmentUsesCustomVirtualHost() {
+		EnvironmentBuilder builder = mock(EnvironmentBuilder.class);
+		RabbitProperties properties = new RabbitProperties();
+		properties.getStream().setVirtualHost("stream-virtual-host");
+		RabbitStreamConfiguration.configure(builder, properties);
+		then(builder).should().virtualHost("stream-virtual-host");
+	}
+
+	@Test
+	void whenStreamVirtualHostIsNotSetButDefaultVirtualHostIsSetThenEnvironmentUsesDefaultVirtualHost() {
+		EnvironmentBuilder builder = mock(EnvironmentBuilder.class);
+		RabbitProperties properties = new RabbitProperties();
+		properties.setVirtualHost("default-virtual-host");
+		RabbitStreamConfiguration.configure(builder, properties);
+		then(builder).should().virtualHost("default-virtual-host");
+	}
+
 	@Test
 	void whenStreamCredentialsAreNotSetThenEnvironmentUsesRabbitCredentials() {
 		EnvironmentBuilder builder = mock(EnvironmentBuilder.class);