From 4139cf497db16fbfba5cbbd105c82fe2190fed86 Mon Sep 17 00:00:00 2001
From: josecoll <jose.coll@r3cev.com>
Date: Tue, 27 Jun 2017 17:22:06 +0100
Subject: [PATCH] Added vault query as tokenizable service. (#926)

* Added vault query as tokenizable service.

* Add JUnit test to test access of vault query service within a flow.

* Improved JUnit test to correctly validate tokenizable behaviour.

* Minor cleanup
---
 .../net/corda/node/internal/AbstractNode.kt   |  2 +-
 .../statemachine/FlowFrameworkTests.kt        | 36 ++++++++++++++++---
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt b/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt
index 18b02e4e53..a0849bd84a 100644
--- a/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt
+++ b/node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt
@@ -475,7 +475,7 @@ abstract class AbstractNode(open val configuration: NodeConfiguration,
         keyManagement = makeKeyManagementService(identity)
         scheduler = NodeSchedulerService(services, database, unfinishedSchedules = busyNodeLatch)
 
-        val tokenizableServices = mutableListOf(storage, network, vault, keyManagement, identity, platformClock, scheduler)
+        val tokenizableServices = mutableListOf(storage, network, vault, vaultQuery, keyManagement, identity, platformClock, scheduler)
         makeAdvertisedServices(tokenizableServices)
         return tokenizableServices
     }
diff --git a/node/src/test/kotlin/net/corda/node/services/statemachine/FlowFrameworkTests.kt b/node/src/test/kotlin/net/corda/node/services/statemachine/FlowFrameworkTests.kt
index 54412d9e67..9e4ea74112 100644
--- a/node/src/test/kotlin/net/corda/node/services/statemachine/FlowFrameworkTests.kt
+++ b/node/src/test/kotlin/net/corda/node/services/statemachine/FlowFrameworkTests.kt
@@ -5,18 +5,18 @@ import co.paralleluniverse.fibers.Suspendable
 import com.google.common.util.concurrent.ListenableFuture
 import net.corda.contracts.asset.Cash
 import net.corda.core.*
+import net.corda.core.contracts.ContractState
 import net.corda.core.contracts.DOLLARS
 import net.corda.core.contracts.DummyState
+import net.corda.core.contracts.StateAndRef
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.generateKeyPair
-import net.corda.core.flows.FlowException
-import net.corda.core.flows.FlowLogic
-import net.corda.core.flows.FlowSessionException
-import net.corda.core.flows.InitiatingFlow
+import net.corda.core.flows.*
 import net.corda.core.identity.Party
 import net.corda.core.messaging.MessageRecipients
 import net.corda.core.node.services.PartyInfo
 import net.corda.core.node.services.ServiceInfo
+import net.corda.core.node.services.queryBy
 import net.corda.core.node.services.unconsumedStates
 import net.corda.core.serialization.OpaqueBytes
 import net.corda.core.serialization.deserialize
@@ -597,6 +597,21 @@ class FlowFrameworkTests {
         assertThatThrownBy { result.getOrThrow() }.hasMessageContaining("Vault").hasMessageContaining("private method")
     }
 
+    @Test
+    fun `verify vault query service is tokenizable by force checkpointing within a flow`() {
+        val ptx = TransactionBuilder(notary = notary1.info.notaryIdentity)
+        ptx.addOutputState(DummyState())
+        val stx = node1.services.signInitialTransaction(ptx)
+
+        node1.registerFlowFactory(VaultQueryFlow::class) {
+            WaitingFlows.Committer(it)
+        }
+        val result = node2.services.startFlow(VaultQueryFlow(stx, node1.info.legalIdentity)).resultFuture
+
+        mockNet.runNetwork()
+        assertThat(result.getOrThrow()).isEmpty()
+    }
+
     @Test
     fun `customised client flow`() {
         val receiveFlowFuture = node2.registerFlowFactory(SendFlow::class) { ReceiveFlow(it) }
@@ -863,6 +878,19 @@ class FlowFrameworkTests {
         }
     }
 
+    @InitiatingFlow
+    private class VaultQueryFlow(val stx: SignedTransaction, val otherParty: Party) : FlowLogic<List<StateAndRef<ContractState>>>() {
+        @Suspendable
+        override fun call(): List<StateAndRef<ContractState>> {
+            send(otherParty, stx)
+            // hold onto reference here to force checkpoint of vaultQueryService and thus
+            // prove it is registered as a tokenizableService in the node
+            val vaultQuerySvc = serviceHub.vaultQueryService
+            waitForLedgerCommit(stx.id)
+            return vaultQuerySvc.queryBy<ContractState>().states
+        }
+    }
+
     @InitiatingFlow(version = 2)
     private class UpgradedFlow(val otherParty: Party) : FlowLogic<Any>() {
         @Suspendable