diff --git a/node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt b/node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt
index 79c4650c61..54c7f933c8 100644
--- a/node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt
+++ b/node/src/test/kotlin/net/corda/node/services/persistence/HibernateConfigurationTest.kt
@@ -10,6 +10,7 @@ import net.corda.core.node.services.Vault
 import net.corda.core.node.services.VaultService
 import net.corda.core.schemas.CommonSchemaV1
 import net.corda.core.schemas.PersistentStateRef
+import net.corda.core.serialization.SerializationDefaults
 import net.corda.core.serialization.deserialize
 import net.corda.core.transactions.SignedTransaction
 import net.corda.finance.DOLLARS
@@ -140,7 +141,7 @@ class HibernateConfigurationTest : TestDependencyInjectionBase() {
 
         // execute query
         val queryResults = entityManager.createQuery(criteriaQuery).resultList
-        val coins = queryResults.map { it.contractState.deserialize<TransactionState<Cash.State>>().data }.sumCash()
+        val coins = queryResults.map { it.contractState.deserialize<TransactionState<Cash.State>>(context = SerializationDefaults.STORAGE_CONTEXT).data }.sumCash()
         assertThat(coins.toDecimal() >= BigDecimal("50.00"))
     }
 
@@ -661,7 +662,7 @@ class HibernateConfigurationTest : TestDependencyInjectionBase() {
         val queryResults = entityManager.createQuery(criteriaQuery).resultList
 
         queryResults.forEach {
-            val contractState = it.contractState.deserialize<TransactionState<ContractState>>()
+            val contractState = it.contractState.deserialize<TransactionState<ContractState>>(context = SerializationDefaults.STORAGE_CONTEXT)
             val cashState = contractState.data as Cash.State
             println("${it.stateRef} with owner: ${cashState.owner.owningKey.toBase58String()}") }
 
@@ -745,7 +746,7 @@ class HibernateConfigurationTest : TestDependencyInjectionBase() {
         // execute query
         val queryResults = entityManager.createQuery(criteriaQuery).resultList
         queryResults.forEach {
-            val contractState = it.contractState.deserialize<TransactionState<ContractState>>()
+            val contractState = it.contractState.deserialize<TransactionState<ContractState>>(context = SerializationDefaults.STORAGE_CONTEXT)
             val cashState = contractState.data as Cash.State
             println("${it.stateRef} with owner ${cashState.owner.owningKey.toBase58String()} and participants ${cashState.participants.map { it.owningKey.toBase58String() }}")
         }
diff --git a/testing/test-utils/src/main/kotlin/net/corda/testing/SerializationTestHelpers.kt b/testing/test-utils/src/main/kotlin/net/corda/testing/SerializationTestHelpers.kt
index 73884f6cab..c3f5872cd7 100644
--- a/testing/test-utils/src/main/kotlin/net/corda/testing/SerializationTestHelpers.kt
+++ b/testing/test-utils/src/main/kotlin/net/corda/testing/SerializationTestHelpers.kt
@@ -68,19 +68,18 @@ fun initialiseTestSerialization() {
         registerScheme(AMQPServerSerializationScheme())
     }
 
-    val AMQP_ENABLE_PROP_NAME = "net.corda.testing.amqp.enable"
-    // TODO: Remove these "if" conditions once we fully switched to AMQP
-    (SerializationDefaults.P2P_CONTEXT as TestSerializationContext).delegate = if (java.lang.Boolean.getBoolean(AMQP_ENABLE_PROP_NAME)) {
-        AMQP_P2P_CONTEXT
-    } else {
-        KRYO_P2P_CONTEXT
-    }
+    (SerializationDefaults.P2P_CONTEXT as TestSerializationContext).delegate = if (isAmqpEnabled()) AMQP_P2P_CONTEXT else KRYO_P2P_CONTEXT
     (SerializationDefaults.RPC_SERVER_CONTEXT as TestSerializationContext).delegate = KRYO_RPC_SERVER_CONTEXT
     (SerializationDefaults.RPC_CLIENT_CONTEXT as TestSerializationContext).delegate = KRYO_RPC_CLIENT_CONTEXT
-    (SerializationDefaults.STORAGE_CONTEXT as TestSerializationContext).delegate = KRYO_STORAGE_CONTEXT
+    (SerializationDefaults.STORAGE_CONTEXT as TestSerializationContext).delegate = if (isAmqpEnabled()) AMQP_STORAGE_CONTEXT else KRYO_STORAGE_CONTEXT
     (SerializationDefaults.CHECKPOINT_CONTEXT as TestSerializationContext).delegate = KRYO_CHECKPOINT_CONTEXT
 }
 
+private const val AMQP_ENABLE_PROP_NAME = "net.corda.testing.amqp.enable"
+
+// TODO: Remove usages of this function when we fully switched to AMQP
+private fun isAmqpEnabled(): Boolean = java.lang.Boolean.getBoolean(AMQP_ENABLE_PROP_NAME)
+
 fun resetTestSerialization() {
     (SerializationDefaults.SERIALIZATION_FACTORY as TestSerializationFactory).delegate = null
     (SerializationDefaults.P2P_CONTEXT as TestSerializationContext).delegate = null