Upgrade to Requery 1.2.1 with Composite Key support (#443)

* Test SELECT WHERE IN composite key using requery 1.2.0

Upgraded Vault Service code to use Requery 1.2.0 SELECT .. WHERE IN

Updated generated schema code with Requery 1.2.0

Upgrade to Requery 1.2.1

Upgrade to Requery 1.2.1 - converted to use update DSL with composite key
Removed redundant JDBC SQL test cases.

Minor updates following PR review comments from RP.

* Streamline companion object initialisation.

build.gradle

@@ -34,7 +34,7 @@ buildscript {
     ext.hibernate_version = '5.2.6.Final'
     ext.h2_version = '1.4.194'
     ext.rxjava_version = '1.2.4'
-    ext.requery_version = '1.1.1'
+    ext.requery_version = '1.2.1'
     ext.dokka_version = '0.9.13'
     ext.crash_version = '1.3.2'
 

core/src/main/kotlin/net/corda/core/node/services/Services.kt

@@ -223,13 +223,13 @@ interface VaultService {
 
     /**
      * Reserve a set of [StateRef] for a given [UUID] unique identifier.
-     * Typically, the unique identifier will refer to a Flow id associated with a [Transaction] in an in-flight flow.
+     * Typically, the unique identifier will refer to a Flow lockId associated with a [Transaction] in an in-flight flow.
      * In the case of coin selection, soft locks are automatically taken upon gathering relevant unconsumed input refs.
      *
      * @throws [StatesNotAvailableException] when not possible to softLock all of requested [StateRef]
      */
     @Throws(StatesNotAvailableException::class)
-    fun softLockReserve(id: UUID, stateRefs: Set<StateRef>)
+    fun softLockReserve(lockId: UUID, stateRefs: Set<StateRef>)
 
     /**
      * Release all or an explicitly specified set of [StateRef] for a given [UUID] unique identifier.
@@ -238,7 +238,7 @@ interface VaultService {
      * In the case of coin selection, softLock are automatically released once previously gathered unconsumed input refs
      * are consumed as part of cash spending.
      */
-    fun softLockRelease(id: UUID, stateRefs: Set<StateRef>? = null)
+    fun softLockRelease(lockId: UUID, stateRefs: Set<StateRef>? = null)
 
     /**
      * Retrieve softLockStates for a given [UUID] or return all softLockStates in vault for a given

node-schemas/generated/source/kapt/main/net/corda/node/services/vault/schemas/Models.java

@@ -9,8 +9,8 @@ import javax.annotation.Generated;
 public class Models {
     public static final EntityModel VAULT = new EntityModelBuilder("vault")
     .addType(VaultStatesEntity.$TYPE)
-    .addType(VaultTxnNoteEntity.$TYPE)
     .addType(VaultCashBalancesEntity.$TYPE)
+    .addType(VaultTxnNoteEntity.$TYPE)
     .build();
 
     private Models() {

node-schemas/generated/source/kapt/main/net/corda/node/services/vault/schemas/VaultCashBalancesEntity.java

@@ -103,6 +103,7 @@ public class VaultCashBalancesEntity implements VaultSchema.VaultCashBalances, P
     .setImmutable(false)
     .setReadOnly(false)
     .setStateless(false)
+    .setView(false)
     .setFactory(new Supplier<VaultCashBalancesEntity>() {
         @Override
         public VaultCashBalancesEntity get() {

node-schemas/generated/source/kapt/main/net/corda/node/services/vault/schemas/VaultStatesEntity.java

@@ -387,6 +387,7 @@ public class VaultStatesEntity implements VaultSchema.VaultStates, Persistable {
     .setImmutable(false)
     .setReadOnly(false)
     .setStateless(false)
+    .setView(false)
     .setFactory(new Supplier<VaultStatesEntity>() {
         @Override
         public VaultStatesEntity get() {

node-schemas/generated/source/kapt/main/net/corda/node/services/vault/schemas/VaultTxnNoteEntity.java

@@ -133,6 +133,7 @@ public class VaultTxnNoteEntity implements VaultSchema.VaultTxnNote, Persistable
     .setImmutable(false)
     .setReadOnly(false)
     .setStateless(false)
+    .setView(false)
     .setFactory(new Supplier<VaultTxnNoteEntity>() {
         @Override
         public VaultTxnNoteEntity get() {

node-schemas/src/test/kotlin/net/corda/node/services/vault/schemas/VaultSchemaTest.kt

@@ -6,6 +6,7 @@ import io.requery.kotlin.`in`
 import io.requery.kotlin.eq
 import io.requery.kotlin.invoke
 import io.requery.kotlin.isNull
+import io.requery.query.RowExpression
 import io.requery.rx.KotlinRxEntityStore
 import io.requery.sql.*
 import io.requery.sql.platform.Generic
@@ -27,8 +28,6 @@ import org.junit.Assert
 import org.junit.Before
 import org.junit.Test
 import rx.Observable
-import java.sql.Connection
-import java.sql.DriverManager
 import java.time.Instant
 import java.util.*
 import java.util.concurrent.CountDownLatch
@@ -48,9 +47,6 @@ class VaultSchemaTest {
 
     var transaction : LedgerTransaction? = null
 
-    var jdbcInstance : Connection? = null
-    val jdbcConn : Connection get() = jdbcInstance!!
-
     @Before
     fun setup() {
         val dataSource = JdbcDataSource()
@@ -62,8 +58,6 @@ class VaultSchemaTest {
         val mode = TableCreationMode.DROP_CREATE
         tables.createTables(mode)
 
-        jdbcInstance = DriverManager.getConnection(dataSource.getURL())
-
         // create dummy test data
         setupDummyData()
     }
@@ -493,6 +487,9 @@ class VaultSchemaTest {
         Assert.assertEquals(3, states.size)
     }
 
+    /**
+     *  Requery composite key tests (using RowExpression introduced in 1.2.1)
+     */
     @Test
     fun testQueryWithCompositeKey() {
         // txn entity with 4 input states (SingleOwnerState x 3, MultiOwnerState x 1)
@@ -500,30 +497,32 @@ class VaultSchemaTest {
         dummyStatesInsert(txn)
 
         data.invoke {
-            // Requery does not support SQL-92 select by composite key:
+            val primaryCompositeKey = listOf(VaultStatesEntity.TX_ID, VaultStatesEntity.INDEX)
-            // Raised Issue:
+            val expression = RowExpression.of(primaryCompositeKey)
-            // https://github.com/requery/requery/issues/434
+            val stateRefs = txn.inputs.map { listOf("'${it.ref.txhash}'", it.ref.index) }
 
-            // Test Requery raw query for single key field
+            val result = select(VaultStatesEntity::class) where (expression.`in`(stateRefs))
-            val refs = txn.inputs.map { it.ref }
+            assertEquals(3, result.get().count())
-            val objArgsTxHash = refs.map { it.txhash.toString() }
+        }
-            val objArgsIndex = refs.map { it.index }
+    }
 
-            val queryByTxHashString = "SELECT * FROM VAULT_STATES WHERE transaction_id IN ?"
+    @Test
-            val resultRawQueryTxHash = raw(VaultStatesEntity::class, queryByTxHashString, *objArgsTxHash.toTypedArray())
+    fun testUpdateWithCompositeKey() {
-            assertEquals(8, resultRawQueryTxHash.count())
+        // txn entity with 4 input states (SingleOwnerState x 3, MultiOwnerState x 1)
+        val txn = createTxnWithTwoStateTypes()
+        dummyStatesInsert(txn)
 
-            val queryByIndexString = "SELECT * FROM VAULT_STATES WHERE output_index IN ?"
+        data.invoke {
-            val resultRawQueryIndex = raw(VaultStatesEntity::class, queryByIndexString, *objArgsIndex.toTypedArray())
+            val primaryCompositeKey = listOf(VaultStatesEntity.TX_ID, VaultStatesEntity.INDEX)
-            assertEquals(18, resultRawQueryIndex.count())
+            val expression = RowExpression.of(primaryCompositeKey)
+            val stateRefs = txn.inputs.map { listOf("'${it.ref.txhash}'", it.ref.index) }
 
-            // Use JDBC native query for composite key
+            val update = update(VaultStatesEntity::class)
-            val stateRefs = refs.fold("") { stateRefs, it -> stateRefs + "('${it.txhash}','${it.index}')," }.dropLast(1)
+                    .set(VaultStatesEntity.LOCK_ID, "")
-            val statement = jdbcConn.createStatement()
+                    .set(VaultStatesEntity.LOCK_UPDATE_TIME, Instant.now())
-            val rs = statement.executeQuery("SELECT transaction_id, output_index, contract_state FROM VAULT_STATES WHERE ((transaction_id, output_index) IN ($stateRefs)) AND (state_status = 0)")
+                    .where (VaultStatesEntity.STATE_STATUS eq Vault.StateStatus.UNCONSUMED)
-            var count = 0
+                    .and (expression.`in`(stateRefs)).get()
-            while (rs.next()) count++
+            assertEquals(3, update.value())
-            assertEquals(3, count)
         }
     }
 
@@ -608,15 +607,16 @@ class VaultSchemaTest {
 
         // release soft lock on states
         data.invoke {
-            val query = select(VaultSchema.VaultStates::class) where (VaultSchema.VaultStates::txId `in` txnIds)
+            val primaryCompositeKey = listOf(VaultStatesEntity.TX_ID, VaultStatesEntity.INDEX)
-                    .and(VaultSchema.VaultStates::lockId eq "LOCK#1")
+            val expression = RowExpression.of(primaryCompositeKey)
-            val result = query.get()
+            val stateRefs = transaction!!.inputs.map { listOf("'${it.ref.txhash}'", it.ref.index) }
-            assertEquals(3, result.count())
+
-            result.forEach {
+            val update = update(VaultStatesEntity::class)
-                it.lockId = ""
+                    .set(VaultStatesEntity.LOCK_ID, "")
-                it.lockUpdateTime = Instant.now()
+                    .set(VaultStatesEntity.LOCK_UPDATE_TIME, Instant.now())
-                upsert(it)
+                    .where (VaultStatesEntity.STATE_STATUS eq Vault.StateStatus.UNCONSUMED)
-            }
+                    .and (expression.`in`(stateRefs)).get()
+            assertEquals(3, update.value())
         }
 
         // select unlocked states
@@ -627,55 +627,6 @@ class VaultSchemaTest {
         }
     }
 
-    @Test
-    fun testMultipleSoftLocksUsingNativeJDBC() {
-        // NOTE:
-        //      - Requery using raw SelectForUpdate not working
-        //      - Requery using raw Update not working
-
-        // using native JDBC
-        val refs = transaction!!.inputs.map { it.ref }
-
-        // insert unconsumed state
-        data.invoke {
-            transaction!!.inputs.forEach {
-                val stateEntity = createStateEntity(it)
-                insert(stateEntity)
-            }
-        }
-
-        // update refs with soft lock id
-        val stateRefs = refs.fold("") { stateRefs, it -> stateRefs + "('${it.txhash}','${it.index}')," }.dropLast(1)
-        val lockId = "LOCK#1"
-        val selectForUpdateStatement = """
-            SELECT transaction_id, output_index, lock_id, lock_timestamp FROM VAULT_STATES
-            WHERE ((transaction_id, output_index) IN ($stateRefs)) FOR UPDATE
-        """
-
-        val statement = jdbcConn.createStatement()
-        val rs = statement.executeQuery(selectForUpdateStatement)
-        while (rs.next()) {
-            val txHash = SecureHash.parse(rs.getString(1))
-            val index = rs.getInt(2)
-            val statement = jdbcConn.createStatement()
-            val updateStatement = """
-             UPDATE VAULT_STATES SET lock_id = '$lockId', lock_timestamp = '${Instant.now()}'
-             WHERE (transaction_id = '$txHash' AND output_index = $index)
-        """
-            statement.executeUpdate(updateStatement)
-        }
-
-        // count locked state refs
-        val selectStatement = """
-            SELECT transaction_id, output_index, contract_state FROM VAULT_STATES
-            WHERE ((transaction_id, output_index) IN ($stateRefs)) AND (lock_id != '')
-        """
-        val rsQuery = statement.executeQuery(selectStatement)
-        var countQuery = 0
-        while (rsQuery.next()) countQuery++
-        assertEquals(3, countQuery)
-    }
-
     @Test
     fun insertWithBigCompositeKey() {
         val keys = (1..314).map { generateKeyPair().public.composite }

node/src/main/kotlin/net/corda/node/services/vault/NodeVaultService.kt

@@ -2,11 +2,13 @@ package net.corda.node.services.vault
 
 import co.paralleluniverse.fibers.Suspendable
 import co.paralleluniverse.strands.Strand
+import io.requery.PersistenceException
 import io.requery.TransactionIsolation
 import io.requery.kotlin.`in`
 import io.requery.kotlin.eq
 import io.requery.kotlin.isNull
 import io.requery.kotlin.notNull
+import io.requery.query.RowExpression
 import net.corda.contracts.asset.Cash
 import net.corda.core.ThreadBox
 import net.corda.core.bufferUntilSubscribed
@@ -32,14 +34,11 @@ import net.corda.core.utilities.trace
 import net.corda.node.services.database.RequeryConfiguration
 import net.corda.node.services.statemachine.FlowStateMachineImpl
 import net.corda.node.services.vault.schemas.*
-import net.corda.node.utilities.StrandLocalTransactionManager
 import net.corda.node.utilities.bufferUntilDatabaseCommit
 import net.corda.node.utilities.wrapWithDatabaseTransaction
-import org.jetbrains.exposed.sql.transactions.TransactionManager
 import rx.Observable
 import rx.subjects.PublishSubject
 import java.security.PublicKey
-import java.sql.Connection
 import java.sql.SQLException
 import java.util.*
 import java.util.concurrent.locks.ReentrantLock
@@ -60,6 +59,9 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
 
     private companion object {
         val log = loggerFor<NodeVaultService>()
+
+        // Define composite primary key used in Requery Expression
+        val stateRefCompositeColumn : RowExpression = RowExpression.of(listOf(VaultStatesEntity.TX_ID, VaultStatesEntity.INDEX))
     }
 
     val configuration = RequeryConfiguration(dataSourceProperties)
@@ -248,83 +250,72 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
     }
 
     @Throws(StatesNotAvailableException::class)
-    override fun softLockReserve(id: UUID, stateRefs: Set<StateRef>) {
+    override fun softLockReserve(lockId: UUID, stateRefs: Set<StateRef>) {
         if (stateRefs.isNotEmpty()) {
-            val stateRefsAsStr = stateRefsToCompositeKeyStr(stateRefs.toList())
             val softLockTimestamp = services.clock.instant()
-            // TODO: awaiting support of UPDATE WHERE <Composite key> IN in Requery DSL
+            val stateRefArgs = stateRefArgs(stateRefs)
-            val updateStatement = """
-                UPDATE VAULT_STATES SET lock_id = '$id', lock_timestamp = '$softLockTimestamp'
-                WHERE ((transaction_id, output_index) IN ($stateRefsAsStr))
-                AND (state_status = 0)
-                AND ((lock_id = '$id') OR (lock_id is null));
-            """
-            val statement = configuration.jdbcSession().createStatement()
-            log.debug(updateStatement)
             try {
-                val rs = statement.executeUpdate(updateStatement)
+                session.withTransaction(TransactionIsolation.REPEATABLE_READ) {
-                if (rs > 0 && rs == stateRefs.size) {
+                    val updatedRows = update(VaultStatesEntity::class)
-                    log.trace("Reserving soft lock states for $id: $stateRefs")
+                            .set(VaultStatesEntity.LOCK_ID, lockId.toString())
-                }
+                            .set(VaultStatesEntity.LOCK_UPDATE_TIME, softLockTimestamp)
-                else {
+                            .where(VaultStatesEntity.STATE_STATUS eq Vault.StateStatus.UNCONSUMED)
-                    // revert partial soft locks
+                            .and((VaultStatesEntity.LOCK_ID eq lockId.toString()) or (VaultStatesEntity.LOCK_ID.isNull()))
-                    val revertUpdateStatement = """
+                            .and(stateRefCompositeColumn.`in`(stateRefArgs)).get().value()
-                        UPDATE VAULT_STATES SET lock_id = null
+                    if (updatedRows > 0 && updatedRows == stateRefs.size) {
-                        WHERE ((transaction_id, output_index) IN ($stateRefsAsStr))
+                        log.trace("Reserving soft lock states for $lockId: $stateRefs")
-                        AND (lock_timestamp = '$softLockTimestamp') AND (lock_id = '$id');
+                    } else {
-                    """
+                        // revert partial soft locks
-                    log.debug(revertUpdateStatement)
+                        val revertUpdatedRows = update(VaultStatesEntity::class)
-                    val rsr = statement.executeUpdate(revertUpdateStatement)
+                                .set(VaultStatesEntity.LOCK_ID, null)
-                    if (rsr > 0) {
+                                .where(VaultStatesEntity.LOCK_UPDATE_TIME eq softLockTimestamp)
-                        log.trace("Reverting $rsr partially soft locked states for $id")
+                                .and(VaultStatesEntity.LOCK_ID eq lockId.toString())
+                                .and(stateRefCompositeColumn.`in`(stateRefArgs)).get().value()
+                        if (revertUpdatedRows > 0) {
+                            log.trace("Reverting $revertUpdatedRows partially soft locked states for $lockId")
+                        }
+                        throw StatesNotAvailableException("Attempted to reserve $stateRefs for $lockId but only $updatedRows rows available")
                     }
-                    throw StatesNotAvailableException("Attempted to reserve $stateRefs for $id but only $rs rows available")
                 }
+            } catch (e: PersistenceException) {
+                log.error("""soft lock update error attempting to reserve states for $lockId and $stateRefs")
+                    $e.
+                """)
+                if (e.cause is StatesNotAvailableException) throw (e.cause as StatesNotAvailableException)
             }
-            catch (e: SQLException) {
-                log.error("""soft lock update error attempting to reserve states: $stateRefs for $id
-                            $e.
-                        """)
-                throw StatesNotAvailableException("Failed to reserve $stateRefs for $id", e)
-            }
-            finally { statement.close() }
         }
     }
 
-    override fun softLockRelease(id: UUID, stateRefs: Set<StateRef>?) {
+    override fun softLockRelease(lockId: UUID, stateRefs: Set<StateRef>?) {
         if (stateRefs == null) {
             session.withTransaction(TransactionIsolation.REPEATABLE_READ) {
                 val update = update(VaultStatesEntity::class)
                         .set(VaultStatesEntity.LOCK_ID, null)
                         .set(VaultStatesEntity.LOCK_UPDATE_TIME, services.clock.instant())
                         .where (VaultStatesEntity.STATE_STATUS eq Vault.StateStatus.UNCONSUMED)
-                        .and (VaultStatesEntity.LOCK_ID eq id.toString()).get()
+                        .and (VaultStatesEntity.LOCK_ID eq lockId.toString()).get()
                 if (update.value() > 0) {
-                    log.trace("Releasing ${update.value()} soft locked states for $id")
+                    log.trace("Releasing ${update.value()} soft locked states for $lockId")
                 }
             }
         }
         else if (stateRefs.isNotEmpty()) {
-            val stateRefsAsStr = stateRefsToCompositeKeyStr(stateRefs.toList())
-            // TODO: awaiting support of UPDATE WHERE <Composite key> IN in Requery DSL
-            val updateStatement = """
-                UPDATE VAULT_STATES SET lock_id = null, lock_timestamp = '${services.clock.instant()}'
-                WHERE (transaction_id, output_index) IN ($stateRefsAsStr)
-                AND (state_status = 0) AND (lock_id = '$id');
-            """
-            val statement = configuration.jdbcSession().createStatement()
-            log.debug(updateStatement)
             try {
-                val rs = statement.executeUpdate(updateStatement)
+                session.withTransaction(TransactionIsolation.REPEATABLE_READ) {
-                if (rs > 0) {
+                    val updatedRows = update(VaultStatesEntity::class)
-                    log.trace("Releasing $rs soft locked states for $id and stateRefs $stateRefs")
+                            .set(VaultStatesEntity.LOCK_ID, null)
+                            .set(VaultStatesEntity.LOCK_UPDATE_TIME, services.clock.instant())
+                            .where(VaultStatesEntity.STATE_STATUS eq Vault.StateStatus.UNCONSUMED)
+                            .and(VaultStatesEntity.LOCK_ID eq lockId.toString())
+                            .and(stateRefCompositeColumn.`in`(stateRefArgs(stateRefs))).get().value()
+                    if (updatedRows > 0) {
+                        log.trace("Releasing $updatedRows soft locked states for $lockId and stateRefs $stateRefs")
+                    }
                 }
-            } catch (e: SQLException) {
+            } catch (e: PersistenceException) {
-                log.error("""soft lock update error attempting to release states for $id and $stateRefs")
+                log.error("""soft lock update error attempting to release states for $lockId and $stateRefs")
-                $e.
+                    $e.
-            """)
+                """)
-            } finally {
-                statement.close()
             }
         }
     }
@@ -568,28 +559,17 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
         // Retrieve all unconsumed states for this transaction's inputs
         val consumedStates = HashSet<StateAndRef<ContractState>>()
         if (tx.inputs.isNotEmpty()) {
-            val stateRefs = stateRefsToCompositeKeyStr(tx.inputs)
+            session.withTransaction(TransactionIsolation.REPEATABLE_READ) {
-            // TODO: using native JDBC until requery supports SELECT WHERE COMPOSITE_KEY IN
+                val result = select(VaultStatesEntity::class).
-            // https://github.com/requery/requery/issues/434
+                             where (stateRefCompositeColumn.`in`(stateRefArgs(tx.inputs))).
-            val statement = configuration.jdbcSession().createStatement()
+                             and (VaultSchema.VaultStates::stateStatus eq Vault.StateStatus.UNCONSUMED)
-            try {
+                result.get().forEach {
-                // TODO: upgrade to Requery 1.2.0 and rewrite with Requery DSL (https://github.com/requery/requery/issues/434)
+                    val txHash = SecureHash.parse(it.txId)
-                val rs = statement.executeQuery("SELECT transaction_id, output_index, contract_state " +
+                    val index = it.index
-                        "FROM vault_states " +
+                    val state = it.contractState.deserialize<TransactionState<ContractState>>(storageKryo())
-                        "WHERE ((transaction_id, output_index) IN ($stateRefs)) " +
-                        "AND (state_status = 0)")
-                while (rs.next()) {
-                    val txHash = SecureHash.parse(rs.getString(1))
-                    val index = rs.getInt(2)
-                    val state = rs.getBytes(3).deserialize<TransactionState<ContractState>>(storageKryo())
                     consumedStates.add(StateAndRef(state, StateRef(txHash, index)))
                 }
-            } catch (e: SQLException) {
-                log.error("""Failed retrieving state refs for: $stateRefs
-                            $e.
-                        """)
             }
-            finally { statement.close() }
         }
 
         // Is transaction irrelevant?
@@ -626,9 +606,9 @@ class NodeVaultService(private val services: ServiceHub, dataSourceProperties: P
     }
 
     /**
-     * Helper method to generate a string formatted list of Composite Keys for SQL IN clause
+     * Helper method to generate a string formatted list of Composite Keys for Requery Expression clause
      */
-    private fun stateRefsToCompositeKeyStr(stateRefs: List<StateRef>): String {
+    private fun stateRefArgs(stateRefs: Iterable<StateRef>): List<List<Any>> {
-        return stateRefs.fold("") { stateRefsAsStr, it -> stateRefsAsStr + "('${it.txhash}','${it.index}')," }.dropLast(1)
+        return stateRefs.map { listOf("'${it.txhash}'", it.index) }
     }
 }

node/src/test/kotlin/net/corda/node/services/vault/NodeVaultServiceTest.kt

@@ -4,15 +4,13 @@ import net.corda.contracts.asset.Cash
 import net.corda.contracts.testing.fillWithSomeTestCash
 import net.corda.core.contracts.*
 import net.corda.core.crypto.composite
-import net.corda.core.flows.FlowException
+import net.corda.core.node.services.StatesNotAvailableException
 import net.corda.core.node.services.TxWritableStorageService
 import net.corda.core.node.services.VaultService
 import net.corda.core.node.services.unconsumedStates
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.utilities.DUMMY_NOTARY
 import net.corda.core.utilities.LogHelper
-import net.corda.node.services.schema.HibernateObserver
-import net.corda.node.services.schema.NodeSchemaService
 import net.corda.node.utilities.configureDatabase
 import net.corda.node.utilities.databaseTransaction
 import net.corda.testing.MEGA_CORP
@@ -132,12 +130,12 @@ class NodeVaultServiceTest {
             assertThat(services.vaultService.softLockedStates<Cash.State>(softLockId)).hasSize(2)
 
             // excluding softlocked states
-            val unlockedStates1 = services.vaultService.unconsumedStates<Cash.State>(includeSoftLockedStates = false)
+            val unlockedStates1 = services.vaultService.unconsumedStates<Cash.State>(includeSoftLockedStates = false).toList()
             assertThat(unlockedStates1).hasSize(1)
 
             // soft lock release one of the states explicitly
             services.vaultService.softLockRelease(softLockId, setOf(unconsumedStates[1].ref))
-            val unlockedStates2 = services.vaultService.unconsumedStates<Cash.State>(includeSoftLockedStates = false)
+            val unlockedStates2 = services.vaultService.unconsumedStates<Cash.State>(includeSoftLockedStates = false).toList()
             assertThat(unlockedStates2).hasSize(2)
 
             // soft lock release the rest by id
@@ -231,7 +229,7 @@ class NodeVaultServiceTest {
 
         // attempt to lock all 3 states with LockId2
         databaseTransaction(database) {
-            assertThatExceptionOfType(FlowException::class.java).isThrownBy(
+            assertThatExceptionOfType(StatesNotAvailableException::class.java).isThrownBy(
                     { vault.softLockReserve(softLockId2, stateRefsToSoftLock) }
             ).withMessageContaining("only 2 rows available").withNoCause()
         }