From 941747d196bdaa835673c8d8c21792e047415baa Mon Sep 17 00:00:00 2001
From: Olga Maciaszek-Sharma <omaciaszeksh@vmware.com>
Date: Mon, 14 Oct 2024 20:38:59 +0200
Subject: [PATCH] Fix tests.

---
 ...oadBalancedRestClientIntegrationTests.java | 11 ++---
 ...edRestTemplateBuilderIntegrationTests.java | 48 ++++++++++++++++---
 2 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestClientIntegrationTests.java b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestClientIntegrationTests.java
index e79988686..eb7d69342 100644
--- a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestClientIntegrationTests.java
+++ b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestClientIntegrationTests.java
@@ -24,7 +24,6 @@
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.cloud.client.ServiceInstance;
-import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.web.client.RestClient;
 
@@ -40,17 +39,17 @@
 		properties = "spring.cloud.loadbalancer.retry.enabled=false")
 public class LoadBalancedRestClientIntegrationTests {
 
-	private final RestClient client;
+	private final RestClient.Builder restClientBuilder;
 
-	@Autowired
-	ApplicationContext context;
 
-	public LoadBalancedRestClientIntegrationTests(@Autowired RestClient.Builder clientBuilder) {
-		this.client = clientBuilder.build();
+	public LoadBalancedRestClientIntegrationTests(@Autowired RestClient.Builder restClientBuilder) {
+		this.restClientBuilder = restClientBuilder;
 	}
 
 	@Test
 	void shouldBuildLoadBalancedRestClientInConstructor() {
+		RestClient client = restClientBuilder.build();
+
 		// Interceptors are not visible in RestClient
 		assertThatThrownBy(() -> client.get().uri("http://test-service").retrieve())
 			.hasMessage("LoadBalancerInterceptor invoked.");
diff --git a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestTemplateBuilderIntegrationTests.java b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestTemplateBuilderIntegrationTests.java
index 239863700..d0dd0655f 100644
--- a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestTemplateBuilderIntegrationTests.java
+++ b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancedRestTemplateBuilderIntegrationTests.java
@@ -16,13 +16,15 @@
 
 package org.springframework.cloud.client.loadbalancer;
 
+import java.net.URI;
+
 import org.junit.jupiter.api.Test;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.web.client.RestTemplateBuilder;
-import org.springframework.context.ApplicationContext;
+import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.context.annotation.Bean;
 import org.springframework.web.client.RestTemplate;
 
@@ -39,9 +41,6 @@ public class LoadBalancedRestTemplateBuilderIntegrationTests {
 
 	private final RestTemplateBuilder restTemplateBuilder;
 
-	@Autowired
-	ApplicationContext context;
-
 	public LoadBalancedRestTemplateBuilderIntegrationTests(@Autowired RestTemplateBuilder restTemplateBuilder) {
 		this.restTemplateBuilder = restTemplateBuilder;
 	}
@@ -51,7 +50,12 @@ void shouldBuildLoadBalancedRestTemplate() {
 		RestTemplate restTemplate = restTemplateBuilder.build();
 
 		assertThat(restTemplate.getInterceptors()).hasSize(1);
-		assertThat(restTemplate.getInterceptors().get(0)).isInstanceOf(BlockingLoadBalancerInterceptor.class);
+		assertThat(restTemplate.getInterceptors()
+				.get(0)).isInstanceOf(DeferringLoadBalancerInterceptor.class);
+		assertThat(((DeferringLoadBalancerInterceptor) restTemplate.getInterceptors()
+				.get(0))
+				.getLoadBalancerInterceptorProvider()
+				.getObject()).isInstanceOf(BlockingLoadBalancerInterceptor.class);
 	}
 
 	@SpringBootConfiguration
@@ -63,6 +67,38 @@ RestTemplateBuilder restTemplateBuilder() {
 			return new RestTemplateBuilder();
 		}
 
+		@Bean
+		LoadBalancerClient testLoadBalancerClient() {
+			return new LoadBalancerClient() {
+				@Override
+				public <T> T execute(String serviceId, LoadBalancerRequest<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public <T> T execute(String serviceId, ServiceInstance serviceInstance,
+						LoadBalancerRequest<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public URI reconstructURI(ServiceInstance instance, URI original) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public ServiceInstance choose(String serviceId) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public <T> ServiceInstance choose(String serviceId, Request<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+			};
+		}
+
+	}
+
 	}
 
-}