diff --git a/core/src/test/kotlin/net/corda/core/flows/ContractUpgradeFlowTest.kt b/core/src/test/kotlin/net/corda/core/flows/ContractUpgradeFlowTest.kt
index bdd3d6ff2f..041a065fb0 100644
--- a/core/src/test/kotlin/net/corda/core/flows/ContractUpgradeFlowTest.kt
+++ b/core/src/test/kotlin/net/corda/core/flows/ContractUpgradeFlowTest.kt
@@ -22,6 +22,7 @@ import net.corda.node.internal.StartedNode
 import net.corda.node.services.Permissions.Companion.startFlow
 import net.corda.testing.contracts.DummyContract
 import net.corda.testing.contracts.DummyContractV2
+import net.corda.testing.contracts.DummyContractV3
 import net.corda.testing.core.ALICE_NAME
 import net.corda.testing.core.BOB_NAME
 import net.corda.testing.core.singleIdentity
@@ -99,17 +100,28 @@ class ContractUpgradeFlowTest {
         mockNet.runNetwork()
 
         val result = resultFuture.resultFuture.getOrThrow()
+        check<DummyContract.State, DummyContractV2.State>(aliceNode, result)
+        check<DummyContract.State, DummyContractV2.State>(bobNode, result)
 
-        fun check(node: StartedNode<MockNode>) {
-            val upgradeTx = node.database.transaction {
-                val wtx = node.services.validatedTransactions.getTransaction(result.ref.txhash)
-                wtx!!.resolveContractUpgradeTransaction(node.services)
-            }
-            assertTrue(upgradeTx.inputs.single().state.data is DummyContract.State)
-            assertTrue(upgradeTx.outputs.single().data is DummyContractV2.State)
+        // Check that the updated state can be further upgraded to V3
+        // Party B authorise the contract state upgrade
+        bobNode.services.startFlow(ContractUpgradeFlow.Authorise(result, DummyContractV3::class.java)).resultFuture.getOrThrow()
+
+        val resultFuture2 = aliceNode.services.startFlow(ContractUpgradeFlow.Initiate(result, DummyContractV3::class.java))
+        mockNet.runNetwork()
+        val result2 = resultFuture2.resultFuture.getOrThrow()
+
+        check<DummyContractV2.State, DummyContractV3.State>(aliceNode, result2)
+        check<DummyContractV2.State, DummyContractV3.State>(bobNode, result2)
+    }
+
+    private inline fun <reified OLD: ContractState, reified NEW: ContractState> check(node: StartedNode<MockNode>, state: StateAndRef<NEW>) {
+        val upgradeTx = node.database.transaction {
+            val wtx = node.services.validatedTransactions.getTransaction(state.ref.txhash)
+            wtx!!.resolveContractUpgradeTransaction(node.services)
         }
-        check(aliceNode)
-        check(bobNode)
+        assertTrue(upgradeTx.inputs.single().state.data is OLD)
+        assertTrue(upgradeTx.outputs.single().data is NEW)
     }
 
     private fun RPCDriverDSL.startProxy(node: StartedNode<MockNode>, user: User): CordaRPCOps {
diff --git a/node/src/main/kotlin/net/corda/node/services/CoreFlowHandlers.kt b/node/src/main/kotlin/net/corda/node/services/CoreFlowHandlers.kt
index bb388eea30..1c89e24228 100644
--- a/node/src/main/kotlin/net/corda/node/services/CoreFlowHandlers.kt
+++ b/node/src/main/kotlin/net/corda/node/services/CoreFlowHandlers.kt
@@ -7,8 +7,8 @@ import net.corda.core.contracts.requireThat
 import net.corda.core.flows.*
 import net.corda.core.identity.Party
 import net.corda.core.internal.ContractUpgradeUtils
-import net.corda.core.transactions.ContractUpgradeWireTransaction
 import net.corda.core.node.StatesToRecord
+import net.corda.core.transactions.ContractUpgradeWireTransaction
 import net.corda.core.transactions.SignedTransaction
 
 // TODO: We should have a whitelist of contracts we're willing to accept at all, and reject if the transaction
@@ -54,9 +54,9 @@ class ContractUpgradeHandler(otherSide: FlowSession) : AbstractStateReplacementF
         // verify outputs matches the proposed upgrade.
         val ourSTX = serviceHub.validatedTransactions.getTransaction(proposal.stateRef.txhash)
         requireNotNull(ourSTX) { "We don't have a copy of the referenced state" }
-        val oldStateAndRef = ourSTX!!.tx.outRef<ContractState>(proposal.stateRef.index)
-        val authorisedUpgrade = serviceHub.contractUpgradeService.getAuthorisedContractUpgrade(oldStateAndRef.ref) ?:
-                throw IllegalStateException("Contract state upgrade is unauthorised. State hash : ${oldStateAndRef.ref}")
+        val oldStateAndRef = ourSTX!!.resolveBaseTransaction(serviceHub).outRef<ContractState>(proposal.stateRef.index)
+        val authorisedUpgrade = serviceHub.contractUpgradeService.getAuthorisedContractUpgrade(oldStateAndRef.ref)
+                ?: throw IllegalStateException("Contract state upgrade is unauthorised. State hash : ${oldStateAndRef.ref}")
         val proposedTx = stx.coreTransaction as ContractUpgradeWireTransaction
         val expectedTx = ContractUpgradeUtils.assembleUpgradeTx(oldStateAndRef, proposal.modification, proposedTx.privacySalt, serviceHub)
         requireThat {
diff --git a/testing/test-utils/src/main/kotlin/net/corda/testing/contracts/DummyContractV3.kt b/testing/test-utils/src/main/kotlin/net/corda/testing/contracts/DummyContractV3.kt
new file mode 100644
index 0000000000..eb8b9a1e97
--- /dev/null
+++ b/testing/test-utils/src/main/kotlin/net/corda/testing/contracts/DummyContractV3.kt
@@ -0,0 +1,36 @@
+package net.corda.testing.contracts
+
+import net.corda.core.contracts.*
+import net.corda.core.identity.AbstractParty
+import net.corda.core.transactions.LedgerTransaction
+
+// The dummy contract doesn't do anything useful. It exists for testing purposes.
+
+/**
+ * Dummy contract state for testing of the upgrade process.
+ */
+class DummyContractV3 : UpgradedContractWithLegacyConstraint<DummyContractV2.State, DummyContractV3.State> {
+    companion object {
+        const val PROGRAM_ID: ContractClassName = "net.corda.testing.contracts.DummyContractV3"
+    }
+
+    override val legacyContract: String = DummyContractV2.PROGRAM_ID
+    override val legacyContractConstraint: AttachmentConstraint = AlwaysAcceptAttachmentConstraint
+
+    data class State(val magicNumber: Int = 0, val owners: List<AbstractParty>) : ContractState {
+        override val participants: List<AbstractParty> = owners
+    }
+
+    interface Commands : CommandData {
+        class Create : TypeOnlyCommandData(), Commands
+        class Move : TypeOnlyCommandData(), Commands
+    }
+
+    override fun upgrade(state: DummyContractV2.State): State {
+        return State(state.magicNumber, state.participants)
+    }
+
+    override fun verify(tx: LedgerTransaction) {
+        // Other verifications.
+    }
+}