CORDA-654 Various MockServices refactorings (#2167)

* Remove MockServices.stateMachineRecordedTransactionMapping which does nothing
* Inline StateLoaderImpl
* Remove unused MockServices
* MockServices well-known identities not needed in a place
* A few things don't need a full-blown ServiceHub

.ci/api-current.txt

@@ -1873,7 +1873,7 @@ public final class net.corda.core.node.services.TimeWindowChecker extends java.l
   @org.jetbrains.annotations.NotNull public final java.time.Clock getClock()
   public final boolean isValid(net.corda.core.contracts.TimeWindow)
 ##
-@net.corda.core.DoNotImplement public interface net.corda.core.node.services.TransactionStorage
+@net.corda.core.DoNotImplement public interface net.corda.core.node.services.TransactionStorage extends net.corda.core.node.StateLoader
   @org.jetbrains.annotations.Nullable public abstract net.corda.core.transactions.SignedTransaction getTransaction(net.corda.core.crypto.SecureHash)
   @org.jetbrains.annotations.NotNull public abstract rx.Observable getUpdates()
   @org.jetbrains.annotations.NotNull public abstract net.corda.core.messaging.DataFeed track()

core/src/main/kotlin/net/corda/core/internal/InternalUtils.kt

@@ -2,9 +2,9 @@
 
 package net.corda.core.internal
 
+import net.corda.core.cordapp.CordappProvider
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.sha256
-import net.corda.core.node.ServiceHub
 import net.corda.core.node.ServicesForResolution
 import net.corda.core.serialization.SerializationContext
 import net.corda.core.transactions.TransactionBuilder
@@ -295,15 +295,15 @@ fun <K, V> Iterable<Pair<K, V>>.toMultiMap(): Map<K, List<V>> = this.groupBy({ i
  * Provide access to internal method for AttachmentClassLoaderTests
  * @suppress
  */
-fun TransactionBuilder.toWireTransaction(services: ServicesForResolution, serializationContext: SerializationContext): WireTransaction {
+fun TransactionBuilder.toWireTransaction(cordappProvider: CordappProvider, serializationContext: SerializationContext): WireTransaction {
-    return toWireTransactionWithContext(services, serializationContext)
+    return toWireTransactionWithContext(cordappProvider, serializationContext)
 }
 
 /**
  * Provide access to internal method for AttachmentClassLoaderTests
  * @suppress
  */
-fun TransactionBuilder.toLedgerTransaction(services: ServiceHub, serializationContext: SerializationContext) = toLedgerTransactionWithContext(services, serializationContext)
+fun TransactionBuilder.toLedgerTransaction(services: ServicesForResolution, serializationContext: SerializationContext) = toLedgerTransactionWithContext(services, serializationContext)
 
 /** Convenience method to get the package name of a class literal. */
 val KClass<*>.packageName: String get() = java.`package`.name

core/src/main/kotlin/net/corda/core/node/ServiceHub.kt

@@ -38,7 +38,9 @@ interface StateLoader {
     // TODO: future implementation to use a Vault state ref -> contract state BLOB table and perform single query bulk load
     // as the existing transaction store will become encrypted at some point
     @Throws(TransactionResolutionException::class)
-    fun loadStates(stateRefs: Set<StateRef>): Set<StateAndRef<ContractState>>
+    fun loadStates(stateRefs: Set<StateRef>): Set<StateAndRef<ContractState>> {
+        return stateRefs.map { StateAndRef(loadState(it), it) }.toSet()
+    }
 }
 
 /**

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

@@ -1,8 +1,12 @@
 package net.corda.core.node.services
 
 import net.corda.core.DoNotImplement
+import net.corda.core.contracts.StateRef
+import net.corda.core.contracts.TransactionResolutionException
+import net.corda.core.contracts.TransactionState
 import net.corda.core.crypto.SecureHash
 import net.corda.core.messaging.DataFeed
+import net.corda.core.node.StateLoader
 import net.corda.core.transactions.SignedTransaction
 import rx.Observable
 
@@ -10,12 +14,18 @@ import rx.Observable
  * Thread-safe storage of transactions.
  */
 @DoNotImplement
-interface TransactionStorage {
+interface TransactionStorage : StateLoader {
     /**
      * Return the transaction with the given [id], or null if no such transaction exists.
      */
     fun getTransaction(id: SecureHash): SignedTransaction?
 
+    @Throws(TransactionResolutionException::class)
+    override fun loadState(stateRef: StateRef): TransactionState<*> {
+        val stx = getTransaction(stateRef.txhash) ?: throw TransactionResolutionException(stateRef.txhash)
+        return stx.resolveBaseTransaction(this).outputs[stateRef.index]
+    }
+
     /**
      * Get a synchronous Observable of updates.  When observations are pushed to the Observer, the vault will already
      * incorporate the update.

core/src/main/kotlin/net/corda/core/transactions/TransactionBuilder.kt

@@ -2,6 +2,7 @@ package net.corda.core.transactions
 
 import co.paralleluniverse.strands.Strand
 import net.corda.core.contracts.*
+import net.corda.core.cordapp.CordappProvider
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.SignableData
 import net.corda.core.crypto.SignatureMetadata
@@ -82,16 +83,16 @@ open class TransactionBuilder(
      * @returns A new [WireTransaction] that will be unaffected by further changes to this [TransactionBuilder].
      */
     @Throws(MissingContractAttachments::class)
-    fun toWireTransaction(services: ServicesForResolution): WireTransaction = toWireTransactionWithContext(services)
+    fun toWireTransaction(services: ServicesForResolution): WireTransaction = toWireTransactionWithContext(services.cordappProvider)
 
-    internal fun toWireTransactionWithContext(services: ServicesForResolution, serializationContext: SerializationContext? = null): WireTransaction {
+    internal fun toWireTransactionWithContext(cordappProvider: CordappProvider, serializationContext: SerializationContext? = null): WireTransaction {
         // Resolves the AutomaticHashConstraints to HashAttachmentConstraints for convenience. The AutomaticHashConstraint
         // allows for less boiler plate when constructing transactions since for the typical case the named contract
         // will be available when building the transaction. In exceptional cases the TransactionStates must be created
         // with an explicit [AttachmentConstraint]
         val resolvedOutputs = outputs.map { state ->
             if (state.constraint is AutomaticHashConstraint) {
-                services.cordappProvider.getContractAttachmentID(state.contract)?.let {
+                cordappProvider.getContractAttachmentID(state.contract)?.let {
                     state.copy(constraint = HashAttachmentConstraint(it))
                 } ?: throw MissingContractAttachments(listOf(state))
             } else {
@@ -106,8 +107,7 @@ open class TransactionBuilder(
     @Throws(AttachmentResolutionException::class, TransactionResolutionException::class)
     fun toLedgerTransaction(services: ServiceHub) = toWireTransaction(services).toLedgerTransaction(services)
 
-    internal fun toLedgerTransactionWithContext(services: ServiceHub, serializationContext: SerializationContext) = toWireTransactionWithContext(services, serializationContext).toLedgerTransaction(services)
+    internal fun toLedgerTransactionWithContext(services: ServicesForResolution, serializationContext: SerializationContext) = toWireTransactionWithContext(services.cordappProvider, serializationContext).toLedgerTransaction(services)
-
     @Throws(AttachmentResolutionException::class, TransactionResolutionException::class, TransactionVerificationException::class)
     fun verify(services: ServiceHub) {
         toLedgerTransaction(services).verify()

node-api/src/test/kotlin/net/corda/nodeapi/internal/AttachmentsClassLoaderStaticContractTests.kt

@@ -1,16 +1,20 @@
 package net.corda.nodeapi.internal
 
+import com.nhaarman.mockito_kotlin.doReturn
+import com.nhaarman.mockito_kotlin.whenever
 import net.corda.core.contracts.*
 import net.corda.core.identity.AbstractParty
 import net.corda.core.identity.Party
+import net.corda.core.node.ServicesForResolution
 import net.corda.core.serialization.deserialize
 import net.corda.core.serialization.serialize
 import net.corda.core.transactions.LedgerTransaction
 import net.corda.core.transactions.TransactionBuilder
+import net.corda.node.internal.cordapp.CordappLoader
+import net.corda.node.internal.cordapp.CordappProviderImpl
 import net.corda.testing.*
-import net.corda.testing.node.MockServices
+import net.corda.testing.node.MockAttachmentStorage
 import org.junit.Assert.*
-import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 
@@ -44,11 +48,8 @@ class AttachmentsClassLoaderStaticContractTests {
         }
     }
 
-    private lateinit var serviceHub: MockServices
+    private val serviceHub = rigorousMock<ServicesForResolution>().also {
-
+        doReturn(CordappProviderImpl(CordappLoader.createWithTestPackages(listOf("net.corda.nodeapi.internal")), MockAttachmentStorage())).whenever(it).cordappProvider
-    @Before
-    fun `create service hub`() {
-        serviceHub = MockServices(cordappPackages = listOf("net.corda.nodeapi.internal"))
     }
 
     @Test

node-api/src/test/kotlin/net/corda/nodeapi/internal/AttachmentsClassLoaderTests.kt

@@ -19,10 +19,8 @@ import net.corda.nodeapi.internal.serialization.attachmentsClassLoaderEnabledPro
 import net.corda.nodeapi.internal.serialization.withTokenContext
 import net.corda.testing.*
 import net.corda.testing.node.MockAttachmentStorage
-import net.corda.testing.node.MockServices
 import org.apache.commons.io.IOUtils
 import org.junit.Assert.*
-import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import java.io.ByteArrayInputStream
@@ -52,14 +50,13 @@ class AttachmentsClassLoaderTests {
     @Rule
     @JvmField
     val testSerialization = SerializationEnvironmentRule()
-    private lateinit var serviceHub: DummyServiceHub
+    private val attachments = MockAttachmentStorage()
-
+    private val cordappProvider = CordappProviderImpl(CordappLoader.createDevMode(listOf(ISOLATED_CONTRACTS_JAR_PATH)), attachments)
-    class DummyServiceHub : MockServices() {
+    private val cordapp get() = cordappProvider.cordapps.first()
-        override val cordappProvider: CordappProviderImpl
+    private val attachmentId get() = cordappProvider.getCordappAttachmentId(cordapp)!!
-                = CordappProviderImpl(CordappLoader.createDevMode(listOf(ISOLATED_CONTRACTS_JAR_PATH)), attachments)
+    private val appContext get() = cordappProvider.getAppContext(cordapp)
-        private val cordapp get() = cordappProvider.cordapps.first()
+    private val serviceHub = rigorousMock<ServiceHub>().also {
-        val attachmentId get() = cordappProvider.getCordappAttachmentId(cordapp)!!
+        doReturn(attachments).whenever(it).attachments
-        val appContext get() = cordappProvider.getAppContext(cordapp)
     }
 
     // These ClassLoaders work together to load 'AnotherDummyContract' in a disposable way, such that even though
@@ -77,12 +74,6 @@ class AttachmentsClassLoaderTests {
     }
 
     class ClassLoaderForTests : URLClassLoader(arrayOf(ISOLATED_CONTRACTS_JAR_PATH), FilteringClassLoader)
-
-    @Before
-    fun `create service hub`() {
-        serviceHub = DummyServiceHub()
-    }
-
     @Test
     fun `dynamically load AnotherDummyContract from isolated contracts jar`() {
         ClassLoaderForTests().use { child ->
@@ -112,8 +103,8 @@ class AttachmentsClassLoaderTests {
 
     @Test
     fun `test MockAttachmentStorage open as jar`() {
-        val storage = serviceHub.attachments
+        val storage = attachments
-        val key = serviceHub.attachmentId
+        val key = attachmentId
         val attachment = storage.openAttachment(key)!!
 
         val jar = attachment.openAsJAR()
@@ -123,9 +114,8 @@ class AttachmentsClassLoaderTests {
 
     @Test
     fun `test overlapping file exception`() {
-        val storage = serviceHub.attachments
+        val storage = attachments
-
+        val att0 = attachmentId
-        val att0 = serviceHub.attachmentId
         val att1 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file.txt", "some data")))
         val att2 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file.txt", "some other data")))
 
@@ -136,9 +126,8 @@ class AttachmentsClassLoaderTests {
 
     @Test
     fun `basic`() {
-        val storage = serviceHub.attachments
+        val storage = attachments
-
+        val att0 = attachmentId
-        val att0 = serviceHub.attachmentId
         val att1 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file1.txt", "some data")))
         val att2 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file2.txt", "some other data")))
 
@@ -169,9 +158,8 @@ class AttachmentsClassLoaderTests {
 
     @Test
     fun `loading class AnotherDummyContract`() {
-        val storage = serviceHub.attachments
+        val storage = attachments
-
+        val att0 = attachmentId
-        val att0 = serviceHub.attachmentId
         val att1 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file1.txt", "some data")))
         val att2 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file2.txt", "some other data")))
 
@@ -194,10 +182,8 @@ class AttachmentsClassLoaderTests {
         val contract = createContract2Cash()
 
         val bytes = contract.serialize()
-
+        val storage = attachments
-        val storage = serviceHub.attachments
+        val att0 = attachmentId
-
-        val att0 = serviceHub.attachmentId
         val att1 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file1.txt", "some data")))
         val att2 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file2.txt", "some other data")))
 
@@ -222,10 +208,8 @@ class AttachmentsClassLoaderTests {
         val context2 = SerializationFactory.defaultFactory.defaultContext.withWhitelisted(data.contract.javaClass)
 
         val bytes = data.serialize(context = context2)
-
+        val storage = attachments
-        val storage = serviceHub.attachments
+        val att0 = attachmentId
-
-        val att0 = serviceHub.attachmentId
         val att1 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file1.txt", "some data")))
         val att2 = storage.importAttachment(ByteArrayInputStream(fakeAttachment("file2.txt", "some other data")))
 
@@ -276,7 +260,7 @@ class AttachmentsClassLoaderTests {
 
     @Test
     fun `test serialization of WireTransaction with dynamically loaded contract`() {
-        val child = serviceHub.appContext.classLoader
+        val child = appContext.classLoader
         val contractClass = Class.forName(ISOLATED_CONTRACT_CLASS_NAME, true, child)
         val contract = contractClass.newInstance() as DummyContractBackdoor
         val tx = contract.generateInitial(MEGA_CORP.ref(0), 42, DUMMY_NOTARY)
@@ -288,7 +272,7 @@ class AttachmentsClassLoaderTests {
                 .withClassLoader(child)
 
         val bytes = run {
-            val wireTransaction = tx.toWireTransaction(serviceHub, context)
+            val wireTransaction = tx.toWireTransaction(cordappProvider, context)
             wireTransaction.serialize(context = context)
         }
         val copiedWireTransaction = bytes.deserialize(context = context)
@@ -307,13 +291,12 @@ class AttachmentsClassLoaderTests {
                 val contractClass = Class.forName(ISOLATED_CONTRACT_CLASS_NAME, true, child)
                 val contract = contractClass.newInstance() as DummyContractBackdoor
                 val tx = contract.generateInitial(MEGA_CORP.ref(0), 42, DUMMY_NOTARY)
-
+                val attachmentRef = attachmentId
-                val attachmentRef = serviceHub.attachmentId
                 val bytes = run {
                     val outboundContext = SerializationFactory.defaultFactory.defaultContext
                             .withServiceHub(serviceHub)
                             .withClassLoader(child)
-                    val wireTransaction = tx.toWireTransaction(serviceHub, outboundContext)
+                    val wireTransaction = tx.toWireTransaction(cordappProvider, outboundContext)
                     wireTransaction.serialize(context = outboundContext)
                 }
                 // use empty attachmentStorage
@@ -340,7 +323,7 @@ class AttachmentsClassLoaderTests {
             val contractClass = Class.forName(ISOLATED_CONTRACT_CLASS_NAME, true, child)
             val contract = contractClass.newInstance() as DummyContractBackdoor
             val outboundContext = SerializationFactory.defaultFactory.defaultContext.withClassLoader(child)
-            val attachmentRef = serviceHub.attachmentId
+            val attachmentRef = attachmentId
             // We currently ignore annotations in attachments, so manually whitelist.
             val inboundContext = SerializationFactory
                     .defaultFactory

node/src/integration-test/kotlin/net/corda/node/services/AttachmentLoadingTests.kt

@@ -1,8 +1,9 @@
 package net.corda.node.services
 
+import com.nhaarman.mockito_kotlin.doReturn
+import com.nhaarman.mockito_kotlin.whenever
 import net.corda.core.contracts.Contract
 import net.corda.core.contracts.PartyAndReference
-import net.corda.core.cordapp.CordappProvider
 import net.corda.core.flows.FlowLogic
 import net.corda.core.flows.UnexpectedFlowEndException
 import net.corda.core.identity.CordaX500Name
@@ -11,6 +12,8 @@ import net.corda.core.internal.concurrent.transpose
 import net.corda.core.internal.createDirectories
 import net.corda.core.internal.div
 import net.corda.core.internal.toLedgerTransaction
+import net.corda.core.node.ServicesForResolution
+import net.corda.core.node.services.IdentityService
 import net.corda.core.serialization.SerializationFactory
 import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.utilities.contextLogger
@@ -23,22 +26,19 @@ import net.corda.testing.DUMMY_NOTARY
 import net.corda.testing.driver.DriverDSLExposedInterface
 import net.corda.testing.driver.NodeHandle
 import net.corda.testing.driver.driver
-import net.corda.testing.node.MockServices
+import net.corda.testing.node.MockAttachmentStorage
 import org.junit.Assert.assertEquals
-import org.junit.Before
 import org.junit.Test
 import java.net.URLClassLoader
 import java.nio.file.Files
 import kotlin.test.assertFailsWith
 
 class AttachmentLoadingTests {
-    private class Services : MockServices() {
+    private val attachments = MockAttachmentStorage()
-        private val provider = CordappProviderImpl(CordappLoader.createDevMode(listOf(isolatedJAR)), attachments)
+    private val provider = CordappProviderImpl(CordappLoader.createDevMode(listOf(isolatedJAR)), attachments)
-        private val cordapp get() = provider.cordapps.first()
+    private val cordapp get() = provider.cordapps.first()
-        val attachmentId get() = provider.getCordappAttachmentId(cordapp)!!
+    private val attachmentId get() = provider.getCordappAttachmentId(cordapp)!!
-        val appContext get() = provider.getAppContext(cordapp)
+    private val appContext get() = provider.getAppContext(cordapp)
-        override val cordappProvider: CordappProvider = provider
-    }
 
     private companion object {
         private val logger = contextLogger()
@@ -70,16 +70,17 @@ class AttachmentLoadingTests {
         }
     }
 
-    private lateinit var services: Services
+    private val services = rigorousMock<ServicesForResolution>().also {
-
+        doReturn(attachments).whenever(it).attachments
-    @Before
+        doReturn(provider).whenever(it).cordappProvider
-    fun setup() {
+        doReturn(rigorousMock<IdentityService>().also {
-        services = Services()
+            doReturn(null).whenever(it).partyFromKey(DUMMY_BANK_A.owningKey)
+        }).whenever(it).identityService
     }
 
     @Test
     fun `test a wire transaction has loaded the correct attachment`() = withTestSerialization {
-        val appClassLoader = services.appContext.classLoader
+        val appClassLoader = appContext.classLoader
         val contractClass = appClassLoader.loadClass(ISOLATED_CONTRACT_ID).asSubclass(Contract::class.java)
         val generateInitialMethod = contractClass.getDeclaredMethod("generateInitial", PartyAndReference::class.java, Integer.TYPE, Party::class.java)
         val contract = contractClass.newInstance()
@@ -89,8 +90,7 @@ class AttachmentLoadingTests {
         contract.verify(ledgerTx)
 
         val actual = ledgerTx.attachments.first()
-        val expected = services.attachments.openAttachment(services.attachmentId)!!
+        val expected = attachments.openAttachment(attachmentId)!!
-
         assertEquals(expected, actual)
     }
 

node/src/main/kotlin/net/corda/node/internal/AbstractNode.kt

@@ -192,8 +192,7 @@ abstract class AbstractNode(val configuration: NodeConfiguration,
         val (startedImpl, schedulerService) = initialiseDatabasePersistence(schemaService, identityService) { database ->
             identityService.loadIdentities(info.legalIdentitiesAndCerts)
             val transactionStorage = makeTransactionStorage(database)
-            val stateLoader = StateLoaderImpl(transactionStorage)
+            val nodeServices = makeServices(keyPairs, schemaService, transactionStorage, database, info, identityService)
-            val nodeServices = makeServices(keyPairs, schemaService, transactionStorage, stateLoader, database, info, identityService)
             val notaryService = makeNotaryService(nodeServices, database)
             val smm = makeStateMachineManager(database)
             val flowStarter = FlowStarterImpl(serverThread, smm)
@@ -201,7 +200,7 @@ abstract class AbstractNode(val configuration: NodeConfiguration,
                     platformClock,
                     database,
                     flowStarter,
-                    stateLoader,
+                    transactionStorage,
                     unfinishedSchedules = busyNodeLatch,
                     serverThread = serverThread)
             if (serverThread is ExecutorService) {
@@ -498,7 +497,7 @@ abstract class AbstractNode(val configuration: NodeConfiguration,
      * Builds node internal, advertised, and plugin services.
      * Returns a list of tokenizable services to be added to the serialisation context.
      */
-    private fun makeServices(keyPairs: Set<KeyPair>, schemaService: SchemaService, transactionStorage: WritableTransactionStorage, stateLoader: StateLoader, database: CordaPersistence, info: NodeInfo, identityService: IdentityService): MutableList<Any> {
+    private fun makeServices(keyPairs: Set<KeyPair>, schemaService: SchemaService, transactionStorage: WritableTransactionStorage, database: CordaPersistence, info: NodeInfo, identityService: IdentityService): MutableList<Any> {
         checkpointStorage = DBCheckpointStorage()
         val metrics = MetricRegistry()
         attachments = NodeAttachmentService(metrics)
@@ -509,7 +508,6 @@ abstract class AbstractNode(val configuration: NodeConfiguration,
                 keyManagementService,
                 schemaService,
                 transactionStorage,
-                stateLoader,
                 MonitoringService(metrics),
                 cordappProvider,
                 database,
@@ -722,18 +720,17 @@ abstract class AbstractNode(val configuration: NodeConfiguration,
             override val keyManagementService: KeyManagementService,
             override val schemaService: SchemaService,
             override val validatedTransactions: WritableTransactionStorage,
-            private val stateLoader: StateLoader,
             override val monitoringService: MonitoringService,
             override val cordappProvider: CordappProviderInternal,
             override val database: CordaPersistence,
             override val myInfo: NodeInfo
-    ) : SingletonSerializeAsToken(), ServiceHubInternal, StateLoader by stateLoader {
+    ) : SingletonSerializeAsToken(), ServiceHubInternal, StateLoader by validatedTransactions {
         override val rpcFlows = ArrayList<Class<out FlowLogic<*>>>()
         override val stateMachineRecordedTransactionMapping = DBTransactionMappingStorage()
         override val auditService = DummyAuditService()
         override val transactionVerifierService by lazy { makeTransactionVerifierService() }
         override val networkMapCache by lazy { NetworkMapCacheImpl(PersistentNetworkMapCache(database), identityService) }
-        override val vaultService by lazy { makeVaultService(keyManagementService, stateLoader, database.hibernateConfig) }
+        override val vaultService by lazy { makeVaultService(keyManagementService, validatedTransactions, database.hibernateConfig) }
         override val contractUpgradeService by lazy { ContractUpgradeServiceImpl() }
         override val attachments: AttachmentStorage get() = this@AbstractNode.attachments
         override val networkService: MessagingService get() = network

node/src/main/kotlin/net/corda/node/internal/StartedNode.kt

@@ -1,14 +1,11 @@
 package net.corda.node.internal
 
-import net.corda.core.contracts.*
 import net.corda.core.flows.FlowLogic
 import net.corda.core.flows.InitiatedBy
 import net.corda.core.internal.VisibleForTesting
 import net.corda.core.messaging.CordaRPCOps
 import net.corda.core.node.NodeInfo
-import net.corda.core.node.StateLoader
 import net.corda.core.node.services.NotaryService
-import net.corda.core.node.services.TransactionStorage
 import net.corda.node.services.api.CheckpointStorage
 import net.corda.node.services.api.StartedNodeServices
 import net.corda.node.services.messaging.MessagingService
@@ -44,17 +41,3 @@ interface StartedNode<out N : AbstractNode> {
         return internals.internalRegisterFlowFactory(smm, initiatingFlowClass, flowFactory, initiatedFlowClass, track)
     }
 }
-
-class StateLoaderImpl(private val validatedTransactions: TransactionStorage) : StateLoader {
-    @Throws(TransactionResolutionException::class)
-    override fun loadState(stateRef: StateRef): TransactionState<*> {
-        val stx = validatedTransactions.getTransaction(stateRef.txhash) ?: throw TransactionResolutionException(stateRef.txhash)
-        return stx.resolveBaseTransaction(this).outputs[stateRef.index]
-    }
-
-    @Throws(TransactionResolutionException::class)
-    // TODO: future implementation to retrieve contract states from a Vault BLOB store
-    override fun loadStates(stateRefs: Set<StateRef>): Set<StateAndRef<ContractState>> {
-        return (stateRefs.map { StateAndRef(loadState(it), it) }).toSet()
-    }
-}

node/src/main/kotlin/net/corda/node/services/persistence/InMemoryStateMachineRecordedTransactionMappingStorage.kt

@@ -1,48 +0,0 @@
-package net.corda.node.services.persistence
-
-import net.corda.core.internal.ThreadBox
-import net.corda.core.internal.bufferUntilSubscribed
-import net.corda.core.crypto.SecureHash
-import net.corda.core.flows.StateMachineRunId
-import net.corda.core.messaging.DataFeed
-import net.corda.core.messaging.StateMachineTransactionMapping
-import net.corda.node.services.api.StateMachineRecordedTransactionMappingStorage
-import rx.subjects.PublishSubject
-import java.util.*
-import javax.annotation.concurrent.ThreadSafe
-
-/**
- * This is a temporary in-memory storage of a state machine id -> txhash mapping
- *
- * TODO persist this instead
- */
-@ThreadSafe
-class InMemoryStateMachineRecordedTransactionMappingStorage : StateMachineRecordedTransactionMappingStorage {
-    private class InnerState {
-        val stateMachineTransactionMap = HashMap<StateMachineRunId, HashSet<SecureHash>>()
-        val updates = PublishSubject.create<StateMachineTransactionMapping>()!!
-    }
-
-    private val mutex = ThreadBox(InnerState())
-
-    override fun addMapping(stateMachineRunId: StateMachineRunId, transactionId: SecureHash) {
-        mutex.locked {
-            stateMachineTransactionMap.getOrPut(stateMachineRunId) { HashSet() }.add(transactionId)
-            updates.onNext(StateMachineTransactionMapping(stateMachineRunId, transactionId))
-        }
-    }
-
-    override fun track():
-            DataFeed<List<StateMachineTransactionMapping>, StateMachineTransactionMapping> {
-        mutex.locked {
-            return DataFeed(
-                    stateMachineTransactionMap.flatMap { entry ->
-                        entry.value.map {
-                            StateMachineTransactionMapping(entry.key, it)
-                        }
-                    },
-                    updates.bufferUntilSubscribed()
-            )
-        }
-    }
-}

node/src/test/kotlin/net/corda/node/services/events/NodeSchedulerServiceTest.kt

@@ -16,7 +16,6 @@ import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.transactions.TransactionBuilder
 import net.corda.core.utilities.days
 import net.corda.node.internal.FlowStarterImpl
-import net.corda.node.internal.StateLoaderImpl
 import net.corda.node.internal.cordapp.CordappLoader
 import net.corda.node.internal.cordapp.CordappProviderImpl
 import net.corda.node.services.api.MonitoringService
@@ -95,7 +94,6 @@ class NodeSchedulerServiceTest : SingletonSerializeAsToken() {
         kms = MockKeyManagementService(identityService, ALICE_KEY)
         val configuration = testNodeConfiguration(Paths.get("."), CordaX500Name("Alice", "London", "GB"))
         val validatedTransactions = MockTransactionStorage()
-        val stateLoader = StateLoaderImpl(validatedTransactions)
         database.transaction {
             services = rigorousMock<Services>().also {
                 doReturn(configuration).whenever(it).configuration
@@ -105,13 +103,13 @@ class NodeSchedulerServiceTest : SingletonSerializeAsToken() {
                 doReturn(myInfo).whenever(it).myInfo
                 doReturn(kms).whenever(it).keyManagementService
                 doReturn(CordappProviderImpl(CordappLoader.createWithTestPackages(listOf("net.corda.testing.contracts")), MockAttachmentStorage())).whenever(it).cordappProvider
-                doReturn(NodeVaultService(testClock, kms, stateLoader, database.hibernateConfig)).whenever(it).vaultService
+                doReturn(NodeVaultService(testClock, kms, validatedTransactions, database.hibernateConfig)).whenever(it).vaultService
                 doReturn(this@NodeSchedulerServiceTest).whenever(it).testReference
 
             }
             smmExecutor = AffinityExecutor.ServiceAffinityExecutor("test", 1)
             mockSMM = StateMachineManagerImpl(services, DBCheckpointStorage(), smmExecutor, database)
-            scheduler = NodeSchedulerService(testClock, database, FlowStarterImpl(smmExecutor, mockSMM), stateLoader, schedulerGatedExecutor, serverThread = smmExecutor)
+            scheduler = NodeSchedulerService(testClock, database, FlowStarterImpl(smmExecutor, mockSMM), validatedTransactions, schedulerGatedExecutor, serverThread = smmExecutor)
             mockSMM.changes.subscribe { change ->
                 if (change is StateMachineManager.Change.Removed && mockSMM.allStateMachines.isEmpty()) {
                     smmHasRemovedAllFlows.countDown()

node/src/test/kotlin/net/corda/node/services/identity/PersistentIdentityServiceTests.kt

@@ -10,13 +10,15 @@ import net.corda.core.internal.cert
 import net.corda.core.internal.toX509CertHolder
 import net.corda.core.node.services.IdentityService
 import net.corda.core.node.services.UnknownAnonymousPartyException
+import net.corda.node.internal.configureDatabase
 import net.corda.nodeapi.internal.crypto.CertificateAndKeyPair
 import net.corda.nodeapi.internal.crypto.CertificateType
 import net.corda.nodeapi.internal.crypto.X509CertificateFactory
 import net.corda.nodeapi.internal.crypto.X509Utilities
 import net.corda.nodeapi.internal.persistence.CordaPersistence
+import net.corda.nodeapi.internal.persistence.DatabaseConfig
 import net.corda.testing.*
-import net.corda.testing.node.MockServices
+import net.corda.testing.node.MockServices.Companion.makeTestDataSourceProperties
 import net.corda.testing.node.makeTestIdentityService
 import org.junit.After
 import org.junit.Before
@@ -30,15 +32,12 @@ import kotlin.test.assertNull
  */
 class PersistentIdentityServiceTests {
     private lateinit var database: CordaPersistence
-    private lateinit var services: MockServices
     private lateinit var identityService: IdentityService
 
     @Before
     fun setup() {
-        val databaseAndServices = MockServices.makeTestDatabaseAndMockServices(emptyList(), PersistentIdentityService(DEV_TRUST_ROOT), initialIdentityName = MEGA_CORP.name)
+        identityService = PersistentIdentityService(DEV_TRUST_ROOT)
-        database = databaseAndServices.first
+        database = configureDatabase(makeTestDataSourceProperties(), DatabaseConfig(), identityService)
-        services = databaseAndServices.second
-        identityService = services.identityService
     }
 
     @After

node/src/test/kotlin/net/corda/node/services/persistence/DBTransactionStorageTests.kt

@@ -5,20 +5,14 @@ import net.corda.core.crypto.Crypto
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.SignatureMetadata
 import net.corda.core.crypto.TransactionSignature
-import net.corda.core.node.StatesToRecord
 import net.corda.core.toFuture
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.transactions.WireTransaction
-import net.corda.node.services.api.VaultServiceInternal
-import net.corda.node.services.schema.HibernateObserver
-import net.corda.node.services.schema.NodeSchemaService
 import net.corda.node.services.transactions.PersistentUniquenessProvider
-import net.corda.node.services.vault.NodeVaultService
 import net.corda.node.internal.configureDatabase
 import net.corda.nodeapi.internal.persistence.CordaPersistence
 import net.corda.nodeapi.internal.persistence.DatabaseConfig
 import net.corda.testing.*
-import net.corda.testing.node.MockServices
 import net.corda.testing.node.MockServices.Companion.makeTestDataSourceProperties
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.After
@@ -35,32 +29,11 @@ class DBTransactionStorageTests {
 
     private lateinit var database: CordaPersistence
     private lateinit var transactionStorage: DBTransactionStorage
-    private lateinit var services: MockServices
-
     @Before
     fun setUp() {
         LogHelper.setLevel(PersistentUniquenessProvider::class)
         val dataSourceProps = makeTestDataSourceProperties()
-        val schemaService = NodeSchemaService()
+        database = configureDatabase(dataSourceProps, DatabaseConfig(), rigorousMock())
-        database = configureDatabase(dataSourceProps, DatabaseConfig(), rigorousMock(), schemaService)
-        database.transaction {
-            services = object : MockServices(BOB_KEY) {
-                override val vaultService: VaultServiceInternal
-                    get() {
-                        val vaultService = NodeVaultService(clock, keyManagementService, stateLoader, database.hibernateConfig)
-                        hibernatePersister = HibernateObserver.install(vaultService.rawUpdates, database.hibernateConfig, schemaService)
-                        return vaultService
-                    }
-
-                override fun recordTransactions(txs: Iterable<SignedTransaction>) {
-                    for (stx in txs) {
-                        validatedTransactions.addTransaction(stx)
-                    }
-                    // Refactored to use notifyAll() as we have no other unit test for that method with multiple transactions.
-                    vaultService.notifyAll(StatesToRecord.ONLY_RELEVANT, txs.map { it.tx })
-                }
-            }
-        }
         newTransactionStorage()
     }
 

testing/node-driver/src/main/kotlin/net/corda/testing/node/MockServices.kt

@@ -4,7 +4,6 @@ import com.google.common.collect.MutableClassToInstanceMap
 import net.corda.core.cordapp.CordappProvider
 import net.corda.core.crypto.*
 import net.corda.core.flows.FlowLogic
-import net.corda.core.flows.StateMachineRunId
 import net.corda.core.identity.CordaX500Name
 import net.corda.core.identity.PartyAndCertificate
 import net.corda.core.messaging.DataFeed
@@ -16,16 +15,13 @@ import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.transactions.SignedTransaction
 import net.corda.node.VersionInfo
-import net.corda.node.internal.StateLoaderImpl
 import net.corda.node.internal.cordapp.CordappLoader
 import net.corda.node.services.api.SchemaService
-import net.corda.node.services.api.StateMachineRecordedTransactionMappingStorage
 import net.corda.node.services.api.VaultServiceInternal
 import net.corda.node.services.api.WritableTransactionStorage
 import net.corda.node.services.identity.InMemoryIdentityService
 import net.corda.node.services.keys.freshCertificate
 import net.corda.node.services.keys.getSigner
-import net.corda.node.services.persistence.InMemoryStateMachineRecordedTransactionMappingStorage
 import net.corda.node.services.schema.HibernateObserver
 import net.corda.node.services.schema.NodeSchemaService
 import net.corda.node.services.transactions.InMemoryTransactionVerifierService
@@ -53,10 +49,9 @@ fun makeTestIdentityService(identities: Iterable<PartyAndCertificate> = emptySet
 open class MockServices(
         cordappLoader: CordappLoader,
         override val validatedTransactions: WritableTransactionStorage,
-        protected val stateLoader: StateLoaderImpl = StateLoaderImpl(validatedTransactions),
         private val initialIdentityName: CordaX500Name = MEGA_CORP.name,
         vararg val keys: KeyPair
-) : ServiceHub, StateLoader by stateLoader {
+) : ServiceHub, StateLoader by validatedTransactions {
     companion object {
         @JvmStatic
         val MOCK_VERSION_INFO = VersionInfo(1, "Mock release", "Mock revision", "Mock Vendor")
@@ -99,9 +94,7 @@ open class MockServices(
                     override val vaultService: VaultServiceInternal = makeVaultService(database.hibernateConfig, schemaService)
 
                     override fun recordTransactions(statesToRecord: StatesToRecord, txs: Iterable<SignedTransaction>) {
-                        for (stx in txs) {
+                        super.recordTransactions(statesToRecord, txs)
-                            validatedTransactions.addTransaction(stx)
-                        }
                         // Refactored to use notifyAll() as we have no other unit test for that method with multiple transactions.
                         vaultService.notifyAll(statesToRecord, txs.map { it.tx })
                     }
@@ -122,15 +115,11 @@ open class MockServices(
 
     override fun recordTransactions(statesToRecord: StatesToRecord, txs: Iterable<SignedTransaction>) {
         txs.forEach {
-            stateMachineRecordedTransactionMapping.addMapping(StateMachineRunId.createRandom(), it.id)
+            validatedTransactions.addTransaction(it)
-        }
-        for (stx in txs) {
-            validatedTransactions.addTransaction(stx)
         }
     }
 
     final override val attachments = MockAttachmentStorage()
-    val stateMachineRecordedTransactionMapping: StateMachineRecordedTransactionMappingStorage = MockStateMachineRecordedTransactionMappingStorage()
     override val identityService: IdentityService = makeTestIdentityService(listOf(MEGA_CORP_IDENTITY, MINI_CORP_IDENTITY, DUMMY_CASH_ISSUER_IDENTITY, DUMMY_NOTARY_IDENTITY))
     override val keyManagementService: KeyManagementService by lazy { MockKeyManagementService(identityService, *keys) }
 
@@ -149,7 +138,7 @@ open class MockServices(
     lateinit var hibernatePersister: HibernateObserver
 
     fun makeVaultService(hibernateConfig: HibernateConfiguration, schemaService: SchemaService): VaultServiceInternal {
-        val vaultService = NodeVaultService(Clock.systemUTC(), keyManagementService, stateLoader, hibernateConfig)
+        val vaultService = NodeVaultService(Clock.systemUTC(), keyManagementService, validatedTransactions, hibernateConfig)
         hibernatePersister = HibernateObserver.install(vaultService.rawUpdates, hibernateConfig, schemaService)
         return vaultService
     }
@@ -201,10 +190,6 @@ class MockKeyManagementService(val identityService: IdentityService,
     }
 }
 
-class MockStateMachineRecordedTransactionMappingStorage(
-        val storage: StateMachineRecordedTransactionMappingStorage = InMemoryStateMachineRecordedTransactionMappingStorage()
-) : StateMachineRecordedTransactionMappingStorage by storage
-
 open class MockTransactionStorage : WritableTransactionStorage, SingletonSerializeAsToken() {
     override fun track(): DataFeed<List<SignedTransaction>, SignedTransaction> {
         return DataFeed(txns.values.toList(), _updatesPublisher)