Made SendKeyForAccountFlow return the created key too.

workflows/src/main/kotlin/com/r3/corda/lib/accounts/workflows/flows/RequestKeyForAccountFlows.kt

@@ -3,7 +3,6 @@ package com.r3.corda.lib.accounts.workflows.flows
 import co.paralleluniverse.fibers.Suspendable
 import com.r3.corda.lib.accounts.contracts.states.AccountInfo
 import com.r3.corda.lib.accounts.workflows.accountService
-import com.r3.corda.lib.accounts.workflows.internal.flows.AccountSearchStatus
 import com.r3.corda.lib.accounts.workflows.internal.flows.createKeyForAccount
 import com.r3.corda.lib.ci.workflows.ProvideKeyFlow
 import com.r3.corda.lib.ci.workflows.RequestKeyFlow
@@ -35,49 +34,38 @@ class RequestKeyForAccountFlow(
         return if (hostSession.counterparty == ourIdentity) {
             serviceHub.createKeyForAccount(accountInfo)
         } else {
-            val accountSearchStatus = hostSession.sendAndReceive<AccountSearchStatus>(accountInfo.identifier.id).unwrap { it }
-            when (accountSearchStatus) {
-                // Maybe the account WAS hosted at this node but has since moved.
-                AccountSearchStatus.NOT_FOUND -> {
-                    throw IllegalStateException("Account Host: ${accountInfo.host} for ${accountInfo.identifier} " +
-                            "(${accountInfo.name}) responded with a not found status - contact them for assistance")
-                }
-                AccountSearchStatus.FOUND -> {
-                    val newKey = subFlow(RequestKeyFlow(hostSession, accountInfo.identifier.id)).owningKey
-                    // Store a local mapping of the account ID to the public key we've just received from the host.
-                    // This allows us to look up the account which the PublicKey is linked to in the future.
-                    // Note that this mapping of KEY -> PARTY persists even when an account moves to another node, the
-                    // assumption being that keys are not moved with the account. If keys DO move with accounts then
-                    // a new API must be added to REPLACE KEY -> PARTY mappings.
-                    //
-                    // The PublicKeyHashToAccountIdMapping table has a primary key constraint over PublicKey, therefore
-                    // a key can only ever be stored once. If you try to store a key twice, then an exception will be
-                    // thrown in respect of the primary key constraint violation.
-                    AnonymousParty(newKey)
-                }
-            }
+            hostSession.send(accountInfo.identifier.id)
+            val newKey = subFlow(RequestKeyFlow(hostSession, accountInfo.identifier.id)).owningKey
+            // Store a local mapping of the account ID to the public key we've just received from the host.
+            // This allows us to look up the account which the PublicKey is linked to in the future.
+            // Note that this mapping of KEY -> PARTY persists even when an account moves to another node, the
+            // assumption being that keys are not moved with the account. If keys DO move with accounts then
+            // a new API must be added to REPLACE KEY -> PARTY mappings.
+            //
+            // The PublicKeyHashToAccountIdMapping table has a primary key constraint over PublicKey, therefore
+            // a key can only ever be stored once. If you try to store a key twice, then an exception will be
+            // thrown in respect of the primary key constraint violation.
+            AnonymousParty(newKey)
         }
     }
 }
 
 /**
  * Responder flow for [RequestKeyForAccountFlow].
  */
-class SendKeyForAccountFlow(val otherSide: FlowSession) : FlowLogic<Unit>() {
+class SendKeyForAccountFlow(private val otherSide: FlowSession) : FlowLogic<AnonymousParty>() {
     @Suspendable
-    override fun call() {
+    override fun call(): AnonymousParty {
         // No need to do anything if the initiating node is us. We can generate a key locally.
         if (otherSide.counterparty == ourIdentity) {
-            return
+            throw FlowException("Should not call on your own")
         }
         val requestedAccountForKey = otherSide.receive(UUID::class.java).unwrap { it }
         val existingAccountInfo = accountService.accountInfo(requestedAccountForKey)
         if (existingAccountInfo == null) {
-            otherSide.send(AccountSearchStatus.NOT_FOUND)
-        } else {
-            otherSide.send(AccountSearchStatus.FOUND)
-            subFlow(ProvideKeyFlow(otherSide))
+            throw FlowException("Account for $requestedAccountForKey not found")
         }
+        return subFlow(ProvideKeyFlow(otherSide))
     }
 }
 
@@ -95,7 +83,7 @@ class RequestKeyForAccount(private val accountInfo: AccountInfo) : FlowLogic<Ano
 }
 
 @InitiatedBy(RequestKeyForAccount::class)
-class SendKeyForAccount(val otherSide: FlowSession) : FlowLogic<Unit>() {
+class SendKeyForAccount(private val otherSide: FlowSession) : FlowLogic<AnonymousParty>() {
     @Suspendable
     override fun call() = subFlow(SendKeyForAccountFlow(otherSide))
 }

workflows/src/main/kotlin/com/r3/corda/lib/accounts/workflows/internal/flows/ConfidentialIdentityUtils.kt

@@ -6,12 +6,6 @@ import net.corda.core.identity.AnonymousParty
 import net.corda.core.node.ServiceHub
 import net.corda.core.serialization.CordaSerializable
 
-@CordaSerializable
-internal enum class AccountSearchStatus {
-    FOUND,
-    NOT_FOUND
-}
-
 @CordaInternal
 fun ServiceHub.createKeyForAccount(accountInfo: AccountInfo) : AnonymousParty {
     val newKey = this.keyManagementService.freshKey(accountInfo.identifier.id)

workflows/src/test/kotlin/com/r3/corda/lib/accounts/workflows/test/RequestKeyForAccountFlowsTest.kt

@@ -1,18 +1,22 @@
 package com.r3.corda.lib.accounts.workflows.test
 
+import com.r3.corda.lib.accounts.contracts.states.AccountInfo
 import com.r3.corda.lib.accounts.workflows.flows.CreateAccount
 import com.r3.corda.lib.accounts.workflows.flows.RequestKeyForAccount
 import com.r3.corda.lib.accounts.workflows.internal.accountService
+import net.corda.core.contracts.UniqueIdentifier
+import net.corda.core.flows.FlowException
 import net.corda.testing.common.internal.testNetworkParameters
 import net.corda.testing.node.MockNetwork
 import net.corda.testing.node.MockNetworkParameters
 import net.corda.testing.node.StartedMockNode
 import net.corda.testing.node.TestCordapp
-import org.hamcrest.CoreMatchers.`is`
 import org.junit.After
-import org.junit.Assert
 import org.junit.Before
 import org.junit.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
 
 class RequestKeyForAccountFlowsTest {
 
@@ -42,6 +46,23 @@ class RequestKeyForAccountFlowsTest {
         network.stopNodes()
     }
 
+    @Test
+    fun `should create locally when request is local`() {
+        val accountA = nodeA.startFlow(CreateAccount("Test_AccountA")).runAndGet(network)
+        val confidentIdentityA = nodeA.startFlow(RequestKeyForAccount(accountA.state.data)).runAndGet(network)
+        assertNotNull(confidentIdentityA)
+        val accountServiceA = nodeA.services.accountService
+        nodeA.transaction {
+            //check if the the key was actually generated by node nodeA
+            nodeA.services.identityService.requireWellKnownPartyFromAnonymous(confidentIdentityA)
+            val keysForAccountA = accountServiceA.accountKeys(accountA.state.data.identifier.id)
+            assertEquals(keysForAccountA, listOf(confidentIdentityA.owningKey))
+        }
+        nodeB.transaction {
+            assertNull(nodeB.services.identityService.partyFromKey(confidentIdentityA.owningKey))
+        }
+    }
+
     //checks if the flow returns nodeA public key for the account
     @Test
     fun `should return key when a node requests`() {
@@ -50,34 +71,36 @@ class RequestKeyForAccountFlowsTest {
         val confidentIdentityA = nodeB.startFlow(RequestKeyForAccount(accountA.state.data)).runAndGet(network)
         val accountServiceA = nodeA.services.accountService
         //verify that nodeA key is returned
-        Assert.assertNotNull(confidentIdentityA)
+        assertNotNull(confidentIdentityA)
         nodeB.transaction {
             //check if the the key was actually generated by node nodeA
             nodeB.services.identityService.requireWellKnownPartyFromAnonymous(confidentIdentityA)
         }
         val keysForAccountA = nodeA.transaction {
             accountServiceA.accountKeys(accountA.state.data.identifier.id)
         }
-        Assert.assertEquals(keysForAccountA, listOf(confidentIdentityA.owningKey))
-
+        assertEquals(keysForAccountA, listOf(confidentIdentityA.owningKey))
     }
 
     //check if it is possible to access account using the public key generated
     @Test
-    fun `should be possible to get the account by newly created key`(){
-        val accountA=nodeA.startFlow(CreateAccount("Test_AccountA")).runAndGet(network)
+    fun `should be possible to get the account by newly created key`() {
+        val accountA = nodeA.startFlow(CreateAccount("Test_AccountA")).runAndGet(network)
         //nodeB request public key
-        val confidentIdentityA=nodeB.startFlow((RequestKeyForAccount(accountA.state.data))).runAndGet(network)
+        val confidentIdentityA = nodeB.startFlow((RequestKeyForAccount(accountA.state.data))).runAndGet(network)
         val accountService = nodeA.services.accountService
-        nodeA.transaction{
+        nodeA.transaction {
             //access the account using accountInfo method ,passing the Public key as parameter
             // and check if the account returned is 'accountA'.
-            Assert.assertThat(accountService.accountInfo(confidentIdentityA.owningKey), `is`(accountA))
+            assertEquals(accountService.accountInfo(confidentIdentityA.owningKey), accountA)
         }
-
-
     }
-}
-
 
+    @Test(expected = FlowException::class)
+    fun `should throw if remote node does not have account`() {
+        val accountA = AccountInfo("Test_Account", nodeA.identity(), UniqueIdentifier())
+        //nodeB request public key
+        nodeB.startFlow((RequestKeyForAccount(accountA))).runAndGet(network)
+    }
 
+}