From 654750ed58d31633faaf6e4c1cd66f506a2383d1 Mon Sep 17 00:00:00 2001
From: Haotian Zhang <928016560@qq.com>
Date: Tue, 1 Aug 2023 20:16:44 +0800
Subject: [PATCH] fix:add deregister judgment. (#1074)

---
 CHANGELOG.md                                  |  1 +
 .../PolarisAutoServiceRegistration.java       | 16 +++++++++++
 .../PolarisAutoServiceRegistrationTest.java   | 27 +++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 22fab005e..3fc40f18b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,3 +23,4 @@
 - [fix:fix gray release examples bug.](https://github.com/Tencent/spring-cloud-tencent/pull/1066)
 - [fix:fix router label feign interceptor order.](https://github.com/Tencent/spring-cloud-tencent/pull/1069)
 - [feat:added polaris weighted round robin load balancer.](https://github.com/Tencent/spring-cloud-tencent/pull/1062)
+- [fix:add deregister judgment.](https://github.com/Tencent/spring-cloud-tencent/pull/1074)
diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistration.java b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistration.java
index c1d4e9203..96b359096 100644
--- a/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistration.java
+++ b/spring-cloud-starter-tencent-polaris-discovery/src/main/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistration.java
@@ -87,6 +87,22 @@ public class PolarisAutoServiceRegistration extends AbstractAutoServiceRegistrat
 		super.registerManagement();
 	}
 
+	@Override
+	protected void deregister() {
+		if (!this.registration.isRegisterEnabled()) {
+			return;
+		}
+		super.deregister();
+	}
+
+	@Override
+	protected void deregisterManagement() {
+		if (!this.registration.isRegisterEnabled()) {
+			return;
+		}
+		super.deregisterManagement();
+	}
+
 	@Override
 	protected Object getConfiguration() {
 		return this.polarisDiscoveryProperties;
diff --git a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistrationTest.java b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistrationTest.java
index 9d78c9033..e1c04a29b 100644
--- a/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistrationTest.java
+++ b/spring-cloud-starter-tencent-polaris-discovery/src/test/java/com/tencent/cloud/polaris/registry/PolarisAutoServiceRegistrationTest.java
@@ -71,6 +71,7 @@ public class PolarisAutoServiceRegistrationTest {
 	@BeforeEach
 	void setUp() {
 		doNothing().when(serviceRegistry).register(nullable(PolarisRegistration.class));
+		doNothing().when(serviceRegistry).deregister(nullable(PolarisRegistration.class));
 
 		polarisAutoServiceRegistration =
 				new PolarisAutoServiceRegistration(serviceRegistry, autoServiceRegistrationProperties, registration,
@@ -111,6 +112,32 @@ public class PolarisAutoServiceRegistrationTest {
 		}).doesNotThrowAnyException();
 	}
 
+	@Test
+	public void testDeregister() {
+		doReturn(false).when(registration).isRegisterEnabled();
+		assertThatCode(() -> {
+			polarisAutoServiceRegistration.registerManagement();
+		}).doesNotThrowAnyException();
+
+		doReturn(true).when(registration).isRegisterEnabled();
+		assertThatCode(() -> {
+			polarisAutoServiceRegistration.deregister();
+		}).doesNotThrowAnyException();
+	}
+
+	@Test
+	public void testDeregisterManagement() {
+		doReturn(false).when(registration).isRegisterEnabled();
+		assertThatCode(() -> {
+			polarisAutoServiceRegistration.registerManagement();
+		}).doesNotThrowAnyException();
+
+		doReturn(true).when(registration).isRegisterEnabled();
+		assertThatCode(() -> {
+			polarisAutoServiceRegistration.deregisterManagement();
+		}).doesNotThrowAnyException();
+	}
+
 	@Test
 	public void testGetAppName() {
 		doReturn("application").when(environment).getProperty(anyString(), anyString());