From d2d5e5feb4a119f9686e61dee3be099b23cf50bb Mon Sep 17 00:00:00 2001
From: DerekYRC <15521077528@163.com>
Date: Mon, 29 Aug 2022 20:42:48 +0800
Subject: [PATCH] Set error handler named EnhancedRestTemplateReporter for
 RestTemplate

---
 .../RpcEnhancementAutoConfiguration.java      | 16 +++++++--
 .../RpcEnhancementAutoConfigurationTest.java  | 35 +------------------
 2 files changed, 14 insertions(+), 37 deletions(-)

diff --git a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfiguration.java b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfiguration.java
index e337f6a6a..85b874009 100644
--- a/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfiguration.java
+++ b/spring-cloud-tencent-rpc-enhancement/src/main/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfiguration.java
@@ -17,6 +17,7 @@
 
 package com.tencent.cloud.rpc.enhancement.config;
 
+import java.util.Collections;
 import java.util.List;
 
 import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
@@ -29,13 +30,14 @@ import com.tencent.cloud.rpc.enhancement.feign.plugin.reporter.SuccessPolarisRep
 import com.tencent.cloud.rpc.enhancement.resttemplate.EnhancedRestTemplateReporter;
 import com.tencent.polaris.api.core.ConsumerAPI;
 
+import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
-import org.springframework.cloud.client.loadbalancer.RestTemplateCustomizer;
+import org.springframework.cloud.client.loadbalancer.LoadBalanced;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.web.client.RestTemplate;
@@ -99,6 +101,10 @@ public class RpcEnhancementAutoConfiguration {
 	@ConditionalOnClass(name = "org.springframework.web.client.RestTemplate")
 	protected static class PolarisRestTemplateAutoConfiguration {
 
+		@LoadBalanced
+		@Autowired(required = false)
+		private List<RestTemplate> restTemplates = Collections.emptyList();
+
 		@Bean
 		public EnhancedRestTemplateReporter polarisRestTemplateResponseErrorHandler(
 				RpcEnhancementReporterProperties properties,
@@ -107,8 +113,12 @@ public class RpcEnhancementAutoConfiguration {
 		}
 
 		@Bean
-		public RestTemplateCustomizer setErrorHandlerCustomizer(EnhancedRestTemplateReporter reporter) {
-			return restTemplate -> restTemplate.setErrorHandler(reporter);
+		public SmartInitializingSingleton setErrorHandlerForRestTemplate(EnhancedRestTemplateReporter reporter) {
+			return () -> {
+				for (RestTemplate restTemplate : restTemplates) {
+					restTemplate.setErrorHandler(reporter);
+				}
+			};
 		}
 	}
 }
diff --git a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java
index 1fa5f826c..7754d8d9e 100644
--- a/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java
+++ b/spring-cloud-tencent-rpc-enhancement/src/test/java/com/tencent/cloud/rpc/enhancement/config/RpcEnhancementAutoConfigurationTest.java
@@ -17,9 +17,6 @@
 
 package com.tencent.cloud.rpc.enhancement.config;
 
-import java.io.IOException;
-import java.net.URI;
-
 import com.tencent.cloud.polaris.context.config.PolarisContextAutoConfiguration;
 import com.tencent.cloud.rpc.enhancement.feign.EnhancedFeignBeanPostProcessor;
 import com.tencent.cloud.rpc.enhancement.feign.EnhancedFeignPluginRunner;
@@ -32,11 +29,7 @@ import org.junit.Test;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.runner.WebApplicationContextRunner;
-import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.cloud.client.loadbalancer.LoadBalanced;
-import org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration;
-import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
-import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.web.client.RestTemplate;
@@ -54,8 +47,7 @@ public class RpcEnhancementAutoConfigurationTest {
 			.withConfiguration(AutoConfigurations.of(
 					PolarisContextAutoConfiguration.class,
 					RpcEnhancementAutoConfiguration.class,
-					PolarisRestTemplateAutoConfigurationTester.class,
-					LoadBalancerAutoConfiguration.class))
+					PolarisRestTemplateAutoConfigurationTester.class))
 			.withPropertyValues("spring.cloud.polaris.circuitbreaker.enabled=true");
 
 	@Test
@@ -82,30 +74,5 @@ public class RpcEnhancementAutoConfigurationTest {
 		RestTemplate restTemplate() {
 			return new RestTemplate();
 		}
-
-		@Bean
-		LoadBalancerClient loadBalancerClient() {
-			return new LoadBalancerClient() {
-				@Override
-				public ServiceInstance choose(String serviceId) {
-					return null;
-				}
-
-				@Override
-				public <T> T execute(String serviceId, LoadBalancerRequest<T> request) throws IOException {
-					return null;
-				}
-
-				@Override
-				public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBalancerRequest<T> request) throws IOException {
-					return null;
-				}
-
-				@Override
-				public URI reconstructURI(ServiceInstance instance, URI original) {
-					return null;
-				}
-			};
-		}
 	}
 }