Make AbstractNode more customisable and upgrade the MockNode implementation.

src/main/kotlin/core/node/AbstractNode.kt

@@ -83,7 +83,8 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
     lateinit var net: MessagingService
 
     open fun start(): AbstractNode {
-        storage = makeStorageService(dir)
+        log.info("Node starting up ...")
+        storage = initialiseStorageService(dir)
         net = makeMessagingService()
         smm = StateMachineManager(services, serverThread)
         wallet = NodeWalletService(services)
@@ -100,13 +101,7 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
         }
         (services.networkMapService as MockNetworkMap).timestampingNodes.add(tsid)
 
-        // We don't have any identity infrastructure right now, so we just throw together the only two identities we
+        identity = makeIdentityService()
-        // know about: our own, and the identity of the remote timestamper node (if any).
-        val knownIdentities = if (timestamperAddress != null)
-            listOf(storage.myLegalIdentity, timestamperAddress.identity)
-        else
-            listOf(storage.myLegalIdentity)
-        identity = FixedIdentityService(knownIdentities)
 
         // This object doesn't need to be referenced from this class because it registers handlers on the network
         // service and so that keeps it from being collected.
@@ -115,6 +110,16 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
         return this
     }
 
+    protected open fun makeIdentityService(): IdentityService {
+        // We don't have any identity infrastructure right now, so we just throw together the only two identities we
+        // know about: our own, and the identity of the remote timestamper node (if any).
+        val knownIdentities = if (timestamperAddress != null)
+            listOf(storage.myLegalIdentity, timestamperAddress.identity)
+        else
+            listOf(storage.myLegalIdentity)
+        return FixedIdentityService(knownIdentities)
+    }
+
     open fun stop() {
         net.stop()
         serverThread.shutdownNow()
@@ -122,9 +127,36 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
 
     protected abstract fun makeMessagingService(): MessagingService
 
-    protected fun makeStorageService(dir: Path): StorageService {
+    protected open fun initialiseStorageService(dir: Path): StorageService {
         val attachments = makeAttachmentStorage(dir)
+        val (identity, keypair) = obtainKeyPair(dir)
+        return constructStorageService(attachments, identity, keypair)
+    }
 
+    protected open fun constructStorageService(attachments: NodeAttachmentStorage, identity: Party, keypair: KeyPair) =
+            StorageServiceImpl(attachments, identity, keypair)
+
+    open inner class StorageServiceImpl(attachments: NodeAttachmentStorage, identity: Party, keypair: KeyPair) : StorageService {
+        protected val tables = HashMap<String, MutableMap<Any, Any>>()
+
+        @Suppress("UNCHECKED_CAST")
+        override fun <K, V> getMap(tableName: String): MutableMap<K, V> {
+            // TODO: This should become a database.
+            synchronized(tables) {
+                return tables.getOrPut(tableName) { Collections.synchronizedMap(HashMap<Any, Any>()) } as MutableMap<K, V>
+            }
+        }
+
+        override val validatedTransactions: MutableMap<SecureHash, SignedTransaction>
+            get() = getMap("validated-transactions")
+
+        override val attachments: AttachmentStorage = attachments
+        override val contractPrograms = contractFactory
+        override val myLegalIdentity = identity
+        override val myLegalIdentityKey = keypair
+    }
+
+    private fun obtainKeyPair(dir: Path): Pair<Party, KeyPair> {
         // Load the private identity key, creating it if necessary. The identity key is a long term well known key that
         // is distributed to other peers and we use it (or a key signed by it) when we need to do something
         // "permissioned". The identity file is what gets distributed and contains the node's legal name along with
@@ -133,7 +165,7 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
         val privKeyFile = dir.resolve(PRIVATE_KEY_FILE_NAME)
         val pubIdentityFile = dir.resolve(PUBLIC_IDENTITY_FILE_NAME)
 
-        val (identity, keypair) = if (!Files.exists(privKeyFile)) {
+        return if (!Files.exists(privKeyFile)) {
             log.info("Identity key not found, generating fresh key!")
             val keypair: KeyPair = generateKeyPair()
             keypair.serialize().writeToFile(privKeyFile)
@@ -154,30 +186,6 @@ abstract class AbstractNode(val dir: Path, val configuration: NodeConfiguration,
             val keypair = Files.readAllBytes(privKeyFile).deserialize<KeyPair>()
             Pair(myIdentity, keypair)
         }
-
-        log.info("Node owned by ${identity.name} starting up ...")
-
-        val ss = object : StorageService {
-            private val tables = HashMap<String, MutableMap<Any, Any>>()
-
-            @Suppress("UNCHECKED_CAST")
-            override fun <K, V> getMap(tableName: String): MutableMap<K, V> {
-                // TODO: This should become a database.
-                synchronized(tables) {
-                    return tables.getOrPut(tableName) { Collections.synchronizedMap(HashMap<Any, Any>()) } as MutableMap<K, V>
-                }
-            }
-
-            override val validatedTransactions: MutableMap<SecureHash, SignedTransaction>
-                get() = getMap("validated-transactions")
-
-            override val attachments: AttachmentStorage = attachments
-            override val contractPrograms = contractFactory
-            override val myLegalIdentity = identity
-            override val myLegalIdentityKey = keypair
-        }
-
-        return ss
     }
 
     private fun makeAttachmentStorage(dir: Path): NodeAttachmentStorage {

src/test/kotlin/core/messaging/AttachmentTests.kt

@@ -90,8 +90,8 @@ class AttachmentTests {
     @Test
     fun maliciousResponse() {
         // Make a node that doesn't do sanity checking at load time.
-        val n0 = network.createNode(null) { path, config, net, ts ->
+        val n0 = network.createNode(null) { path, config, mock, ts ->
-            object : MockNetwork.MockNode(path, config, net, ts) {
+            object : MockNetwork.MockNode(path, config, mock, ts) {
                 override fun start(): MockNetwork.MockNode {
                     super.start()
                     (storage.attachments as NodeAttachmentStorage).checkAttachmentsOnLoad = false

src/test/kotlin/core/node/MockNode.kt

@@ -19,15 +19,18 @@ package core.node
 import com.google.common.jimfs.Configuration
 import com.google.common.jimfs.Jimfs
 import com.google.common.util.concurrent.MoreExecutors
+import core.Party
 import core.messaging.InMemoryMessagingNetwork
 import core.messaging.LegallyIdentifiableNode
 import core.messaging.MessagingService
+import core.testutils.TEST_KEYS_TO_CORP_MAP
 import core.utilities.loggerFor
 import org.slf4j.Logger
 import java.nio.file.Files
 import java.nio.file.Path
 import java.util.*
 import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
 
 /**
  * A mock node brings up a suite of in-memory services in a fast manner suitable for unit testing.
@@ -37,11 +40,13 @@ import java.util.concurrent.ExecutorService
  * for message exchanges to take place (and associated handlers to run), you must call the [runNetwork]
  * method.
  */
-class MockNetwork {
+class MockNetwork(private val threadPerNode: Boolean = false) {
     private var counter = 0
     val filesystem = Jimfs.newFileSystem(Configuration.unix())
     val messagingNetwork = InMemoryMessagingNetwork()
 
+    val identities = ArrayList<Party>(TEST_KEYS_TO_CORP_MAP.values)
+
     private val _nodes = ArrayList<MockNode>()
     /** A read only view of the current set of executing nodes. */
     val nodes: List<MockNode> = _nodes
@@ -50,34 +55,49 @@ class MockNetwork {
         Files.createDirectory(filesystem.getPath("/nodes"))
     }
 
-    open class MockNode(dir: Path, config: NodeConfiguration, val network: InMemoryMessagingNetwork,
+    open class MockNode(dir: Path, config: NodeConfiguration, val mockNet: MockNetwork,
-                        withTimestamper: LegallyIdentifiableNode?) : AbstractNode(dir, config, withTimestamper) {
+                        withTimestamper: LegallyIdentifiableNode?, val forcedID: Int = -1) : AbstractNode(dir, config, withTimestamper) {
         override val log: Logger = loggerFor<MockNode>()
-        override val serverThread: ExecutorService = MoreExecutors.newDirectExecutorService()
+        override val serverThread: ExecutorService =
+                if (mockNet.threadPerNode)
+                    Executors.newSingleThreadExecutor()
+                else
+                    MoreExecutors.newDirectExecutorService()
 
         // We only need to override the messaging service here, as currently everything that hits disk does so
         // through the java.nio API which we are already mocking via Jimfs.
 
         override fun makeMessagingService(): MessagingService {
-            return network.createNode(true).second.start().get()
+            if (forcedID == -1)
+                return mockNet.messagingNetwork.createNode(!mockNet.threadPerNode).second.start().get()
+            else
+                return mockNet.messagingNetwork.createNodeWithID(!mockNet.threadPerNode, forcedID).start().get()
         }
 
+        override fun makeIdentityService() = FixedIdentityService(mockNet.identities)
+
         override fun start(): MockNode {
             super.start()
+            mockNet.identities.add(storage.myLegalIdentity)
             return this
         }
     }
 
-    fun createNode(withTimestamper: LegallyIdentifiableNode?,
+    /** Returns a started node, optionally created by the passed factory method */
-                   factory: (Path, NodeConfiguration, network: InMemoryMessagingNetwork, LegallyIdentifiableNode?) -> MockNode = { p, n, n2, l -> MockNode(p, n, n2, l) }): MockNode {
+    fun createNode(withTimestamper: LegallyIdentifiableNode?, forcedID: Int = -1,
-        val path = filesystem.getPath("/nodes/$counter")
+                   factory: ((Path, NodeConfiguration, network: MockNetwork, LegallyIdentifiableNode?) -> MockNode)? = null): MockNode {
-        Files.createDirectory(path)
+        val newNode = forcedID == -1
+        val id = if (newNode) counter++ else forcedID
+
+        val path = filesystem.getPath("/nodes/$id")
+        if (newNode)
+            Files.createDirectories(path.resolve("attachments"))
         val config = object : NodeConfiguration {
-            override val myLegalName: String = "Mock Company $counter"
+            override val myLegalName: String = "Mock Company $id"
         }
-        val node = factory(path, config, messagingNetwork, withTimestamper).start()
+        val fac = factory ?: { p, n, n2, l -> MockNode(p, n, n2, l, id) }
+        val node = fac(path, config, this, withTimestamper).start()
         _nodes.add(node)
-        counter++
         return node
     }
 
@@ -98,8 +118,8 @@ class MockNetwork {
     /**
      * Sets up a two node network in which the first node runs a timestamping service and the other doesn't.
      */
-    fun createTwoNodes(): Pair<MockNode, MockNode> {
+    fun createTwoNodes(factory: ((Path, NodeConfiguration, network: MockNetwork, LegallyIdentifiableNode?) -> MockNode)? = null): Pair<MockNode, MockNode> {
         require(nodes.isEmpty())
-        return Pair(createNode(null), createNode(nodes[0].legallyIdentifableAddress))
+        return Pair(createNode(null, -1, factory), createNode(nodes[0].legallyIdentifableAddress, -1, factory))
     }
 }