From 818f8a5f77183d265e43337e78a80d85e685f6c1 Mon Sep 17 00:00:00 2001
From: Michal Kit <michal.kit@r3.com>
Date: Mon, 24 Jul 2017 10:30:43 +0100
Subject: [PATCH] Addressing review comments

---
 .../core/node/services/NetworkMapCache.kt     |  2 +-
 .../node/services/AdvertisedServiceTests.kt   | 43 +++++++++++++------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/core/src/main/kotlin/net/corda/core/node/services/NetworkMapCache.kt b/core/src/main/kotlin/net/corda/core/node/services/NetworkMapCache.kt
index e5821d9ef2..9fc9f521fe 100644
--- a/core/src/main/kotlin/net/corda/core/node/services/NetworkMapCache.kt
+++ b/core/src/main/kotlin/net/corda/core/node/services/NetworkMapCache.kt
@@ -123,7 +123,7 @@ interface NetworkMapCache {
      * Returns a service identity advertised by one of the nodes on the network
      * @param type Specifies the type of the service
      */
-    fun getServiceOf(type: ServiceType): Party? {
+    fun getAnyServiceOfType(type: ServiceType): Party? {
         for (node in partyNodes) {
             val serviceIdentities = node.serviceIdentities(type)
             if (serviceIdentities.isNotEmpty()) {
diff --git a/node/src/integration-test/kotlin/net/corda/node/services/AdvertisedServiceTests.kt b/node/src/integration-test/kotlin/net/corda/node/services/AdvertisedServiceTests.kt
index 2e66e4e016..f64850c8e0 100644
--- a/node/src/integration-test/kotlin/net/corda/node/services/AdvertisedServiceTests.kt
+++ b/node/src/integration-test/kotlin/net/corda/node/services/AdvertisedServiceTests.kt
@@ -1,27 +1,42 @@
 package net.corda.node.services
 
-import com.google.common.util.concurrent.Futures
-import net.corda.core.getOrThrow
+import co.paralleluniverse.fibers.Suspendable
+import net.corda.core.flows.FlowLogic
+import net.corda.core.flows.StartableByRPC
+import net.corda.core.messaging.startFlow
 import net.corda.core.node.services.ServiceInfo
 import net.corda.core.node.services.ServiceType
-import net.corda.testing.DUMMY_BANK_A
-import net.corda.testing.DUMMY_BANK_B
-import net.corda.testing.node.NodeBasedTest
+import net.corda.nodeapi.User
+import net.corda.testing.driver.driver
 import org.bouncycastle.asn1.x500.X500Name
 import org.junit.Test
-import kotlin.test.assertEquals
+import kotlin.test.assertTrue
 
-class AdvertisedServiceTests : NodeBasedTest() {
+class AdvertisedServiceTests {
     private val serviceName = X500Name("CN=Custom Service,O=R3,OU=corda,L=London,C=GB")
     private val serviceType = ServiceType.corda.getSubType("custom")
+    private val user = "bankA"
+    private val pass = "passA"
+
+
+    @StartableByRPC
+    class ServiceTypeCheckingFlow : FlowLogic<Boolean>() {
+
+        @Suspendable
+        override fun call(): Boolean {
+            return serviceHub.networkMapCache.getAnyServiceOfType(ServiceType.corda.getSubType("custom")) != null;
+        }
+    }
 
     @Test
-    fun `service is accessible through getServiceOf`() {
-        val (bankA) = Futures.allAsList(
-                startNode(DUMMY_BANK_A.name),
-                startNode(DUMMY_BANK_B.name, advertisedServices = setOf(ServiceInfo(serviceType, serviceName)))
-        ).getOrThrow()
-        val serviceParty = bankA.services.networkMapCache.getServiceOf(serviceType)
-        assertEquals(serviceName, serviceParty?.name)
+    fun `service is accessible through getAnyServiceOfType`() {
+        driver(startNodesInProcess = true) {
+            val bankA = startNode(rpcUsers = listOf(User(user, pass, setOf(startFlowPermission<ServiceTypeCheckingFlow>())))).get()
+            val bankB = startNode(advertisedServices = setOf(ServiceInfo(serviceType, serviceName))).get()
+            bankA.rpcClientToNode().use(user, pass) { connection ->
+                val result = connection.proxy.startFlow(::ServiceTypeCheckingFlow).returnValue.get()
+                assertTrue(result)
+            }
+        }
     }
 }