From 94f50ed37b10271aac1e17d5a9dc6c531a57c15a Mon Sep 17 00:00:00 2001
From: Spencer Gibb <spencer@gibb.us>
Date: Wed, 30 Jan 2019 16:36:04 -0500
Subject: [PATCH] Changes test to not rely on nuber of calls for correct state.

---
 ...outeDefinitionLocatorIntegrationTests.java | 71 +++++++++++--------
 1 file changed, 41 insertions(+), 30 deletions(-)

diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/discovery/DiscoveryClientRouteDefinitionLocatorIntegrationTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/discovery/DiscoveryClientRouteDefinitionLocatorIntegrationTests.java
index 8ba0ba4c..4a46f4de 100644
--- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/discovery/DiscoveryClientRouteDefinitionLocatorIntegrationTests.java
+++ b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/discovery/DiscoveryClientRouteDefinitionLocatorIntegrationTests.java
@@ -20,7 +20,7 @@ package org.springframework.cloud.gateway.discovery;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -53,6 +53,9 @@ public class DiscoveryClientRouteDefinitionLocatorIntegrationTests {
     @Autowired
 	private ApplicationEventPublisher publisher;
 
+    @Autowired
+	private TestDiscoveryClient discoveryClient;
+
     @Test
     public void newServiceAddsRoute() {
 		List<Route> routes = routeLocator.getRoutes()
@@ -60,6 +63,8 @@ public class DiscoveryClientRouteDefinitionLocatorIntegrationTests {
 				.collectList().block();
 		assertThat(routes).hasSize(1);
 
+		discoveryClient.multiple();
+
 		publisher.publishEvent(new HeartbeatEvent(this, 1L));
 
 		routes = routeLocator.getRoutes()
@@ -73,39 +78,45 @@ public class DiscoveryClientRouteDefinitionLocatorIntegrationTests {
 	protected static class Config {
 
 		@Bean
-		DiscoveryClient discoveryClient() {
-			DefaultServiceInstance instance1 = new DefaultServiceInstance("service1", "localhost", 8001,
-					false);
-			DefaultServiceInstance instance2 = new DefaultServiceInstance("service2", "localhost", 8001,
-					false);
-			return new DiscoveryClient() {
+		TestDiscoveryClient discoveryClient() {
+			return new TestDiscoveryClient();
+		}
+	}
 
-				AtomicInteger calls = new AtomicInteger(0);
+	private static class TestDiscoveryClient implements DiscoveryClient {
+		AtomicBoolean single = new AtomicBoolean(true);
 
-				@Override
-				public String description() {
-					return null;
-				}
+		DefaultServiceInstance instance1 = new DefaultServiceInstance("service1_1",
+				"service1", "localhost", 8001, false);
+		DefaultServiceInstance instance2 = new DefaultServiceInstance("service2_1",
+				"service2", "localhost", 8001, false);
 
-				@Override
-				public List<ServiceInstance> getInstances(String serviceId) {
-					if (serviceId.equals("service1")) {
-						return Collections.singletonList(instance1);
-					}
-					if (serviceId.equals("service2")) {
-						return Collections.singletonList(instance2);
-					}
-					return Collections.emptyList();
-				}
+		public void multiple() {
+			single.set(false);
+		}
 
-				@Override
-				public List<String> getServices() {
-					if (calls.compareAndSet(0, 1)) {
-						return Collections.singletonList("service1");
-					}
-					return Arrays.asList("service1", "service2");
-				}
-			};
+		@Override
+		public String description() {
+			return null;
+		}
+
+		@Override
+		public List<ServiceInstance> getInstances(String serviceId) {
+			if (serviceId.equals("service1")) {
+				return Collections.singletonList(instance1);
+			}
+			if (serviceId.equals("service2")) {
+				return Collections.singletonList(instance2);
+			}
+			return Collections.emptyList();
+		}
+
+		@Override
+		public List<String> getServices() {
+			if (single.get()) {
+				return Collections.singletonList("service1");
+			}
+			return Arrays.asList("service1", "service2");
 		}
 	}
 }