From cca71e3d6eae39664f649cbe13a49fe9dd4953a0 Mon Sep 17 00:00:00 2001
From: Tommy Lillehagen <tommy.lillehagen@r3.com>
Date: Mon, 12 Feb 2018 10:18:18 +0000
Subject: [PATCH] Make RPC scaffolding available for Node

---
 .../main/kotlin/net/corda/behave/node/Node.kt | 22 ++++++++++++++++
 .../corda/behave/scenarios/ScenarioState.kt   | 26 ++++---------------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/experimental/behave/src/main/kotlin/net/corda/behave/node/Node.kt b/experimental/behave/src/main/kotlin/net/corda/behave/node/Node.kt
index e5c9f69bb2..ed6b9f097a 100644
--- a/experimental/behave/src/main/kotlin/net/corda/behave/node/Node.kt
+++ b/experimental/behave/src/main/kotlin/net/corda/behave/node/Node.kt
@@ -9,10 +9,15 @@ import net.corda.behave.logging.getLogger
 import net.corda.behave.monitoring.PatternWatch
 import net.corda.behave.node.configuration.*
 import net.corda.behave.process.JarCommand
+import net.corda.behave.seconds
 import net.corda.behave.service.Service
 import net.corda.behave.service.ServiceSettings
 import net.corda.behave.ssh.MonitoringSSHClient
 import net.corda.behave.ssh.SSHClient
+import net.corda.client.rpc.CordaRPCClient
+import net.corda.client.rpc.CordaRPCClientConfiguration
+import net.corda.core.messaging.CordaRPCOps
+import net.corda.core.utilities.NetworkHostAndPort
 import org.apache.commons.io.FileUtils
 import java.io.File
 import java.time.Duration
@@ -146,6 +151,23 @@ class Node(
         }).start()
     }
 
+    fun <T> rpc(action: (CordaRPCOps) -> T): T {
+        var result: T? = null
+        val user = config.users.first()
+        val address = config.nodeInterface
+        val targetHost = NetworkHostAndPort(address.host, address.rpcPort)
+        val config = CordaRPCClientConfiguration(
+                connectionMaxRetryInterval = 10.seconds
+        )
+        log.info("Establishing RPC connection to ${targetHost.host} on port ${targetHost.port} ...")
+        CordaRPCClient(targetHost, config).use(user.username, user.password) {
+            log.info("RPC connection to ${targetHost.host}:${targetHost.port} established")
+            val client = it.proxy
+            result = action(client)
+        }
+        return result ?: error("Failed to run RPC action")
+    }
+
     override fun toString(): String {
         return "Node(name = ${config.name}, version = ${config.distribution.version})"
     }
diff --git a/experimental/behave/src/scenario/kotlin/net/corda/behave/scenarios/ScenarioState.kt b/experimental/behave/src/scenario/kotlin/net/corda/behave/scenarios/ScenarioState.kt
index 02e2af9c22..f6cfb32298 100644
--- a/experimental/behave/src/scenario/kotlin/net/corda/behave/scenarios/ScenarioState.kt
+++ b/experimental/behave/src/scenario/kotlin/net/corda/behave/scenarios/ScenarioState.kt
@@ -3,11 +3,7 @@ package net.corda.behave.scenarios
 import net.corda.behave.logging.getLogger
 import net.corda.behave.network.Network
 import net.corda.behave.node.Node
-import net.corda.behave.seconds
-import net.corda.client.rpc.CordaRPCClient
-import net.corda.client.rpc.CordaRPCClientConfiguration
 import net.corda.core.messaging.CordaRPCOps
-import net.corda.core.utilities.NetworkHostAndPort
 import org.assertj.core.api.Assertions.assertThat
 
 class ScenarioState {
@@ -54,29 +50,17 @@ class ScenarioState {
         assertThat(network?.waitUntilRunning()).isTrue()
     }
 
-    fun withNetwork(action: ScenarioState.() -> Unit) {
+    inline fun <T> withNetwork(action: ScenarioState.() -> T): T {
         ensureNetworkIsRunning()
-        action()
+        return action()
     }
 
-    fun <T> withClient(nodeName: String, action: (CordaRPCOps) -> T): T {
-        var result: T? = null
+    inline fun <T> withClient(nodeName: String, crossinline action: (CordaRPCOps) -> T): T {
         withNetwork {
-            val node = node(nodeName)
-            val user = node.config.users.first()
-            val address = node.config.nodeInterface
-            val targetHost = NetworkHostAndPort(address.host, address.rpcPort)
-            val config = CordaRPCClientConfiguration(
-                    connectionMaxRetryInterval = 10.seconds
-            )
-            log.info("Establishing RPC connection to ${targetHost.host} on port ${targetHost.port} ...")
-            CordaRPCClient(targetHost, config).use(user.username, user.password) {
-                log.info("RPC connection to ${targetHost.host}:${targetHost.port} established")
-                val client = it.proxy
-                result = action(client)
+            return node(nodeName).rpc {
+                action(it)
             }
         }
-        return result ?: error("Failed to run RPC action")
     }
 
     fun stopNetwork() {