diff --git a/core/src/main/kotlin/net/corda/core/flows/TransactionKeyFlow.kt b/core/src/main/kotlin/net/corda/core/flows/TransactionKeyFlow.kt
index b1567f0557..58ff5837df 100644
--- a/core/src/main/kotlin/net/corda/core/flows/TransactionKeyFlow.kt
+++ b/core/src/main/kotlin/net/corda/core/flows/TransactionKeyFlow.kt
@@ -31,7 +31,6 @@ class TransactionKeyFlow(val otherSide: Party,
     override fun call(): LinkedHashMap<Party, AnonymisedIdentity> {
         progressTracker.currentStep = AWAITING_KEY
         val legalIdentityAnonymous = serviceHub.keyManagementService.freshKeyAndCert(serviceHub.myInfo.legalIdentityAndCert, revocationEnabled)
-        serviceHub.identityService.registerAnonymousIdentity(legalIdentityAnonymous.identity, serviceHub.myInfo.legalIdentity, legalIdentityAnonymous.certPath)
 
         // Special case that if we're both parties, a single identity is generated
         val identities = LinkedHashMap<Party, AnonymisedIdentity>()
@@ -39,6 +38,7 @@ class TransactionKeyFlow(val otherSide: Party,
             identities.put(otherSide, legalIdentityAnonymous)
         } else {
             val otherSideAnonymous = sendAndReceive<AnonymisedIdentity>(otherSide, legalIdentityAnonymous).unwrap { validateIdentity(otherSide, it) }
+            serviceHub.identityService.registerAnonymousIdentity(otherSideAnonymous.identity, otherSide, otherSideAnonymous.certPath)
             identities.put(serviceHub.myInfo.legalIdentity, legalIdentityAnonymous)
             identities.put(otherSide, otherSideAnonymous)
         }
diff --git a/docs/source/example-code/src/integration-test/kotlin/net/corda/docs/IntegrationTestingTutorial.kt b/docs/source/example-code/src/integration-test/kotlin/net/corda/docs/IntegrationTestingTutorial.kt
index c51e115f09..e0e6f6be4a 100644
--- a/docs/source/example-code/src/integration-test/kotlin/net/corda/docs/IntegrationTestingTutorial.kt
+++ b/docs/source/example-code/src/integration-test/kotlin/net/corda/docs/IntegrationTestingTutorial.kt
@@ -8,6 +8,7 @@ import net.corda.core.getOrThrow
 import net.corda.core.messaging.startFlow
 import net.corda.core.node.services.ServiceInfo
 import net.corda.core.node.services.Vault
+import net.corda.core.node.services.vault.QueryCriteria
 import net.corda.core.utilities.OpaqueBytes
 import net.corda.testing.ALICE
 import net.corda.testing.BOB
@@ -57,8 +58,9 @@ class IntegrationTestingTutorial {
             // END 2
 
             // START 3
-            val bobVaultUpdates = bobProxy.vaultAndUpdates().second
-            val aliceVaultUpdates = aliceProxy.vaultAndUpdates().second
+            val criteria = QueryCriteria.VaultQueryCriteria(status = Vault.StateStatus.ALL)
+            val (_, bobVaultUpdates) = bobProxy.vaultTrackByCriteria<Cash.State>(Cash.State::class.java, criteria)
+            val (_, aliceVaultUpdates) = aliceProxy.vaultTrackByCriteria<Cash.State>(Cash.State::class.java, criteria)
             // END 3
 
             // START 4
@@ -70,8 +72,7 @@ class IntegrationTestingTutorial {
                             i.DOLLARS,
                             issueRef,
                             bob.nodeInfo.legalIdentity,
-                            notary.nodeInfo.notaryIdentity,
-                            false // Not anonymised
+                            notary.nodeInfo.notaryIdentity
                     ).returnValue)
                 }
             }.forEach(Thread::join) // Ensure the stack of futures is populated.
@@ -94,7 +95,7 @@ class IntegrationTestingTutorial {
 
             // START 5
             for (i in 1..10) {
-                bobProxy.startFlow(::CashPaymentFlow, i.DOLLARS, alice.nodeInfo.legalIdentity, false).returnValue.getOrThrow()
+                bobProxy.startFlow(::CashPaymentFlow, i.DOLLARS, alice.nodeInfo.legalIdentity).returnValue.getOrThrow()
             }
 
             aliceVaultUpdates.expectEvents {