Improved support for testing custom schemas using a MockNetwork. (#1450)

* Improved support for testing custom schemas using a MockNetwork.

* Removed `requiredSchemas` from CordaPluginREgistry configuration
Custom schema registration now uses classpath scanning (CorDapps) and explicit registration (tests) following same mechanisms as flow registration.

* Updated following PR review feedback.

* Helper function to return Kotlin object instance fixed and moved to core InternalUtils class.

* Fixed auto-scanning Unit test to assert correct registration of custom schema.

* cleanup comment.

* Changes following rebase from master.

client/rpc/src/integration-test/java/net/corda/client/rpc/CordaRPCJavaClientTest.java

@@ -10,6 +10,7 @@ import net.corda.core.utilities.OpaqueBytes;
 import net.corda.finance.flows.AbstractCashFlow;
 import net.corda.finance.flows.CashIssueFlow;
 import net.corda.finance.flows.CashPaymentFlow;
+import net.corda.finance.schemas.*;
 import net.corda.node.internal.Node;
 import net.corda.node.services.transactions.ValidatingNotaryService;
 import net.corda.nodeapi.User;
@@ -52,6 +53,7 @@ public class CordaRPCJavaClientTest extends NodeBasedTest {
         Set<ServiceInfo> services = new HashSet<>(singletonList(new ServiceInfo(ValidatingNotaryService.Companion.getType(), null)));
         CordaFuture<Node> nodeFuture = startNode(getALICE().getName(), 1, services, singletonList(rpcUser), emptyMap());
         node = nodeFuture.get();
+        node.registerCustomSchemas(Collections.singleton(CashSchemaV1.INSTANCE));
         client = new CordaRPCClient(requireNonNull(node.getConfiguration().getRpcAddress()), null, getDefault(), false);
     }
 

client/rpc/src/integration-test/kotlin/net/corda/client/rpc/CordaRPCClientTest.kt

@@ -13,6 +13,7 @@ import net.corda.finance.contracts.getCashBalances
 import net.corda.finance.flows.CashException
 import net.corda.finance.flows.CashIssueFlow
 import net.corda.finance.flows.CashPaymentFlow
+import net.corda.finance.schemas.CashSchemaV1
 import net.corda.node.internal.Node
 import net.corda.node.services.FlowPermissions.Companion.startFlowPermission
 import net.corda.node.services.transactions.ValidatingNotaryService
@@ -44,6 +45,7 @@ class CordaRPCClientTest : NodeBasedTest() {
     @Before
     fun setUp() {
         node = startNode(ALICE.name, rpcUsers = listOf(rpcUser), advertisedServices = setOf(ServiceInfo(ValidatingNotaryService.type))).getOrThrow()
+        node.registerCustomSchemas(setOf(CashSchemaV1))
         client = CordaRPCClient(node.configuration.rpcAddress!!, initialiseSerialization = false)
     }
 

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

@@ -26,6 +26,7 @@ import java.util.zip.Deflater
 import java.util.zip.ZipEntry
 import java.util.zip.ZipOutputStream
 import kotlin.reflect.KClass
+import kotlin.reflect.full.createInstance
 
 val Throwable.rootCause: Throwable get() = cause?.rootCause ?: this
 fun Throwable.getStackTraceAsString() = StringWriter().also { printStackTrace(PrintWriter(it)) }.toString()
@@ -239,6 +240,11 @@ fun <T> Any.declaredField(name: String): DeclaredField<T> = DeclaredField(javaCl
  */
 fun <T> Any.declaredField(clazz: KClass<*>, name: String): DeclaredField<T> = DeclaredField(clazz.java, name, this)
 
+/** creates a new instance if not a Kotlin object */
+fun <T: Any> KClass<T>.objectOrNewInstance(): T {
+    return this.objectInstance ?: this.createInstance()
+}
+
 /**
  * A simple wrapper around a [Field] object providing type safe read and write access using [value], ignoring the field's
  * visibility.

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

@@ -22,13 +22,4 @@ abstract class CordaPluginRegistry {
      * @return true if you register types, otherwise you will be filtered out of the list of plugins considered in future.
      */
     open fun customizeSerialization(custom: SerializationCustomization): Boolean = false
-
-    /**
-     * Optionally, custom schemas to be used for contract state persistence and vault custom querying
-     *
-     * For example, if you implement the [QueryableState] interface on a new [ContractState]
-     * it needs to be registered here if you wish to perform custom queries on schema entity attributes using the
-     * [VaultQueryService] API
-     */
-    open val requiredSchemas: Set<MappedSchema> get() = emptySet()
 }

docs/source/api-persistence.rst

@@ -74,6 +74,18 @@ other ``MappedSchema``.
    ``QueryableState`` being persisted. This will change in due course. Similarly, it does not currently support
    configuring ``SchemaOptions`` but will do so in the future.
 
+Custom schema registration
+--------------------------
+Custom contract schemas are automatically registered at startup time for CorDapps. The node bootstrap process will scan
+for schemas (any class that extends the ``MappedSchema`` interface) in the `plugins` configuration directory in your CorDapp jar.
+
+For testing purposes it is necessary to manually register custom schemas as follows:
+
+- Tests using ``MockNetwork`` and ``MockNode`` must explicitly register custom schemas using the `registerCustomSchemas()` method of ``MockNode``
+- Tests using ``MockServices`` must explicitly register schemas using `customSchemas` attribute of the ``MockServices`` `makeTestDatabaseAndMockServices()` helper method.
+
+.. note:: Tests using the `DriverDSL` will automatically register your custom schemas if they are in the same project structure as the driver call.
+
 Object relational mapping
 -------------------------
 The persisted representation of a ``QueryableState`` should be an instance of a ``PersistentState`` subclass,

docs/source/api-vault-query.rst

@@ -69,7 +69,8 @@ There are four implementations of this interface which can be chained together t
 
 4. ``VaultCustomQueryCriteria`` provides the means to specify one or many arbitrary expressions on attributes defined by a custom contract state that implements its own schema as described in the :doc:`Persistence </api-persistence>` documentation and associated examples. Custom criteria expressions are expressed using one of several type-safe ``CriteriaExpression``: BinaryLogical, Not, ColumnPredicateExpression, AggregateFunctionExpression. The ``ColumnPredicateExpression`` allows for specification arbitrary criteria using the previously enumerated operator types. The ``AggregateFunctionExpression`` allows for the specification of an aggregate function type (sum, avg, max, min, count) with optional grouping and sorting. Furthermore, a rich DSL is provided to enable simple construction of custom criteria using any combination of ``ColumnPredicate``. See the ``Builder`` object in ``QueryCriteriaUtils`` for a complete specification of the DSL. 
 
-    .. note:: It is a requirement to register any custom contract schemas to be used in Vault Custom queries in the associated `CordaPluginRegistry` configuration for the respective CorDapp using the ``requiredSchemas`` configuration field (which specifies a set of `MappedSchema`)
+    .. note:: custom contract schemas are automatically registered upon node startup for CorDapps. Please refer to
+    :doc:`Persistence </api-persistence>` for mechanisms of registering custom schemas for different testing purposes.
 
 All ``QueryCriteria`` implementations are composable using ``and`` and ``or`` operators.
 

docs/source/changelog.rst

@@ -12,6 +12,10 @@ UNRELEASED
 * About half of the code in test-utils has been moved to a new module ``node-driver``,
   and the test scope modules are now located in a ``testing`` directory.
 
+* Removed `requireSchemas` CordaPluginRegistry configuration item.
+  Custom schemas are now automatically located using classpath scanning for deployed CorDapps.
+  Improved support for testing custom schemas in MockNode and MockServices using explicit registration.
+
 * Contract Upgrades: deprecated RPC authorisation / deauthorisation API calls in favour of equivalent flows in ContractUpgradeFlow.
   Implemented contract upgrade persistence using JDBC backed persistent map.
 

docs/source/example-code/src/test/kotlin/net/corda/docs/CustomVaultQueryTest.kt

@@ -7,6 +7,7 @@ import net.corda.core.utilities.getOrThrow
 import net.corda.finance.*
 import net.corda.finance.contracts.getCashBalances
 import net.corda.finance.flows.CashIssueFlow
+import net.corda.finance.schemas.CashSchemaV1
 import net.corda.node.services.network.NetworkMapService
 import net.corda.node.services.transactions.ValidatingNotaryService
 import net.corda.testing.DUMMY_NOTARY
@@ -38,6 +39,8 @@ class CustomVaultQueryTest {
 
         nodeA.registerInitiatedFlow(TopupIssuerFlow.TopupIssuer::class.java)
         nodeA.installCordaService(CustomVaultQuery.Service::class.java)
+        nodeA.registerCustomSchemas(setOf(CashSchemaV1))
+        nodeB.registerCustomSchemas(setOf(CashSchemaV1))
     }
 
     @After

docs/source/example-code/src/test/kotlin/net/corda/docs/FxTransactionBuildTutorialTest.kt

@@ -7,6 +7,7 @@ import net.corda.core.utilities.getOrThrow
 import net.corda.finance.*
 import net.corda.finance.contracts.getCashBalances
 import net.corda.finance.flows.CashIssueFlow
+import net.corda.finance.schemas.CashSchemaV1
 import net.corda.node.services.network.NetworkMapService
 import net.corda.node.services.transactions.ValidatingNotaryService
 import net.corda.testing.DUMMY_NOTARY
@@ -33,6 +34,8 @@ class FxTransactionBuildTutorialTest {
                 advertisedServices = *arrayOf(ServiceInfo(NetworkMapService.type), notaryService))
         nodeA = mockNet.createPartyNode(notaryNode.network.myAddress)
         nodeB = mockNet.createPartyNode(notaryNode.network.myAddress)
+        nodeA.registerCustomSchemas(setOf(CashSchemaV1))
+        nodeB.registerCustomSchemas(setOf(CashSchemaV1))
         nodeB.registerInitiatedFlow(ForeignExchangeRemoteFlow::class.java)
     }
 

docs/source/writing-cordapps.rst

@@ -78,9 +78,6 @@ The ``CordaPluginRegistry`` class defines the following:
 * ``customizeSerialization``, which can be overridden to provide a list of the classes to be whitelisted for object
   serialisation, over and above those tagged with the ``@CordaSerializable`` annotation. See :doc:`serialization`
 
-* ``requiredSchemas``, which can be overridden to return a set of the MappedSchemas to use for persistence and vault
-  queries
-
 The ``WebServerPluginRegistry`` class defines the following:
 
 * ``webApis``, which can be overridden to return a list of JAX-RS annotated REST access classes. These classes will be

finance/src/main/kotlin/net/corda/finance/plugin/FinancePluginRegistry.kt

@@ -1,11 +0,0 @@
-package net.corda.finance.plugin
-
-import net.corda.core.node.CordaPluginRegistry
-import net.corda.core.schemas.MappedSchema
-import net.corda.finance.schemas.CashSchemaV1
-
-class FinancePluginRegistry : CordaPluginRegistry() {
-    override val requiredSchemas: Set<MappedSchema> = setOf(
-            CashSchemaV1
-    )
-}

finance/src/main/resources/META-INF/services/net.corda.core.node.CordaPluginRegistry

@@ -1 +0,0 @@
-net.corda.finance.plugin.FinancePluginRegistry

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

@@ -24,6 +24,7 @@ import net.corda.core.node.PluginServiceHub
 import net.corda.core.node.ServiceEntry
 import net.corda.core.node.services.*
 import net.corda.core.node.services.NetworkMapCache.MapChange
+import net.corda.core.schemas.MappedSchema
 import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.transactions.SignedTransaction
@@ -73,7 +74,6 @@ import rx.Observable
 import java.io.IOException
 import java.lang.reflect.InvocationTargetException
 import java.nio.file.Path
-import java.nio.file.Paths
 import java.security.KeyPair
 import java.security.KeyStoreException
 import java.security.cert.CertificateFactory
@@ -214,6 +214,7 @@ abstract class AbstractNode(open val configuration: NodeConfiguration,
             installCordaServices()
             registerCordappFlows()
             _services.rpcFlows += cordappLoader.findRPCFlows()
+            registerCustomSchemas(cordappLoader.findCustomSchemas())
 
             runOnStop += network::stop
         }
@@ -655,7 +656,7 @@ abstract class AbstractNode(open val configuration: NodeConfiguration,
         override val monitoringService = MonitoringService(MetricRegistry())
         override val validatedTransactions = makeTransactionStorage()
         override val transactionVerifierService by lazy { makeTransactionVerifierService() }
-        override val schemaService by lazy { NodeSchemaService(pluginRegistries.flatMap { it.requiredSchemas }.toSet()) }
+        override val schemaService by lazy { NodeSchemaService() }
         override val networkMapCache by lazy { PersistentNetworkMapCache(this) }
         override val vaultService by lazy { NodeVaultService(this) }
         override val contractUpgradeService by lazy { ContractUpgradeServiceImpl() }
@@ -703,4 +704,8 @@ abstract class AbstractNode(open val configuration: NodeConfiguration,
         override fun jdbcSession(): Connection = database.createSession()
     }
 
+    fun registerCustomSchemas(schemas: Set<MappedSchema>) {
+        database.hibernateConfig.schemaService.registerCustomSchemas(schemas)
+    }
+
 }

node/src/main/kotlin/net/corda/node/internal/classloading/CordappLoader.kt

@@ -6,13 +6,11 @@ import net.corda.core.flows.ContractUpgradeFlow
 import net.corda.core.flows.FlowLogic
 import net.corda.core.flows.InitiatedBy
 import net.corda.core.flows.StartableByRPC
-import net.corda.core.internal.div
-import net.corda.core.internal.exists
-import net.corda.core.internal.isRegularFile
-import net.corda.core.internal.list
+import net.corda.core.internal.*
 import net.corda.core.node.NodeInfo
 import net.corda.core.node.services.CordaService
 import net.corda.core.node.services.ServiceType
+import net.corda.core.schemas.MappedSchema
 import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.utilities.debug
 import net.corda.core.utilities.loggerFor
@@ -121,6 +119,10 @@ class CordappLoader private constructor (val cordappClassPath: List<Path>) {
         return found + coreFlows
     }
 
+    fun findCustomSchemas(): Set<MappedSchema> {
+        return scanResult?.getClassesWithSuperclass(MappedSchema::class)?.toSet() ?: emptySet()
+    }
+
     private fun scanCordapps(): ScanResult? {
         logger.info("Scanning CorDapps in $cordappClassPath")
         return if (cordappClassPath.isNotEmpty())
@@ -144,21 +146,28 @@ class CordappLoader private constructor (val cordappClassPath: List<Path>) {
         }
     }
 
-    private fun <T : Any> ScanResult.getClassesWithAnnotation(type: KClass<T>, annotation: KClass<out Annotation>): List<Class<out T>> {
-        fun loadClass(className: String): Class<out T>? {
-            return try {
-                appClassLoader.loadClass(className) as Class<T>
-            } catch (e: ClassCastException) {
-                logger.warn("As $className is annotated with ${annotation.qualifiedName} it must be a sub-type of ${type.java.name}")
-                null
-            } catch (e: Exception) {
-                logger.warn("Unable to load class $className", e)
-                null
-            }
+    private fun <T : Any> loadClass(className: String, type: KClass<T>): Class<out T>? {
+        return try {
+            appClassLoader.loadClass(className) as Class<T>
+        } catch (e: ClassCastException) {
+            logger.warn("As $className must be a sub-type of ${type.java.name}")
+            null
+        } catch (e: Exception) {
+            logger.warn("Unable to load class $className", e)
+            null
         }
+    }
 
+    private fun <T : Any> ScanResult.getClassesWithSuperclass(type: KClass<T>): List<T> {
+        return getNamesOfSubclassesOf(type.java)
+                .mapNotNull { loadClass(it, type) }
+                .filterNot { Modifier.isAbstract(it.modifiers) }
+                .map { it.kotlin.objectOrNewInstance() }
+    }
+
+    private fun <T : Any> ScanResult.getClassesWithAnnotation(type: KClass<T>, annotation: KClass<out Annotation>): List<Class<out T>> {
         return getNamesOfClassesWithAnnotation(annotation.java)
-                .mapNotNull { loadClass(it) }
+                .mapNotNull { loadClass(it, type) }
                 .filterNot { Modifier.isAbstract(it.modifiers) }
     }
 }

node/src/main/kotlin/net/corda/node/services/api/SchemaService.kt

@@ -30,5 +30,11 @@ interface SchemaService {
      * or via custom logic in this service.
      */
     fun generateMappedObject(state: ContractState, schema: MappedSchema): PersistentState
+
+    /**
+     * Registration mechanism to add custom contract schemas that extend the [MappedSchema] class.
+     */
+    fun registerCustomSchemas(customSchemas: Set<MappedSchema>)
+
 }
 //DOCEND SchemaService

node/src/main/kotlin/net/corda/node/services/schema/NodeSchemaService.kt

@@ -61,13 +61,13 @@ class NodeSchemaService(customSchemas: Set<MappedSchema> = emptySet()) : SchemaS
 
     // Required schemas are those used by internal Corda services
     // For example, cash is used by the vault for coin selection (but will be extracted as a standalone CorDapp in future)
-    val requiredSchemas: Map<MappedSchema, SchemaService.SchemaOptions> =
+    private val requiredSchemas: Map<MappedSchema, SchemaService.SchemaOptions> =
             mapOf(Pair(CommonSchemaV1, SchemaService.SchemaOptions()),
                   Pair(VaultSchemaV1, SchemaService.SchemaOptions()),
                   Pair(NodeInfoSchemaV1, SchemaService.SchemaOptions()),
                   Pair(NodeServicesV1, SchemaService.SchemaOptions()))
 
-    override val schemaOptions: Map<MappedSchema, SchemaService.SchemaOptions> = requiredSchemas.plus(customSchemas.map {
+    override var schemaOptions: Map<MappedSchema, SchemaService.SchemaOptions> = requiredSchemas.plus(customSchemas.map {
         mappedSchema -> Pair(mappedSchema, SchemaService.SchemaOptions())
     })
 
@@ -92,4 +92,10 @@ class NodeSchemaService(customSchemas: Set<MappedSchema> = emptySet()) : SchemaS
             return VaultSchemaV1.VaultFungibleStates(state.owner, state.amount.quantity, state.amount.token.issuer.party, state.amount.token.issuer.reference, state.participants)
         return (state as QueryableState).generateMappedObject(schema)
     }
+
+    override fun registerCustomSchemas(_customSchemas: Set<MappedSchema>) {
+        schemaOptions = schemaOptions.plus(_customSchemas.map {
+            mappedSchema -> Pair(mappedSchema, SchemaService.SchemaOptions())
+        })
+    }
 }

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

@@ -321,9 +321,8 @@ class HibernateQueryCriteriaParser(val contractType: Class<out ContractState>,
             e.message?.let { message ->
                 if (message.contains("Not an entity"))
                     throw VaultQueryException("""
-                    Please register the entity '${entityClass.name.substringBefore('$')}' class in your CorDapp's CordaPluginRegistry configuration (requiredSchemas attribute)
-                    and ensure you have declared (in supportedSchemas()) and mapped (in generateMappedObject()) the schema in the associated contract state's QueryableState interface implementation.
-                    See https://docs.corda.net/persistence.html?highlight=persistence for more information""")
+                    Please register the entity '${entityClass.name.substringBefore('$')}'
+                    See https://docs.corda.net/api-persistence.html#custom-schema-registration for more information""")
             }
             throw VaultQueryException("Parsing error: ${e.message}")
         }

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

@@ -28,14 +28,14 @@ import java.lang.Exception
 import java.util.*
 import javax.persistence.Tuple
 
-class HibernateVaultQueryImpl(hibernateConfig: HibernateConfiguration,
+class HibernateVaultQueryImpl(val hibernateConfig: HibernateConfiguration,
                               val vault: VaultService) : SingletonSerializeAsToken(), VaultQueryService {
     companion object {
         val log = loggerFor<HibernateVaultQueryImpl>()
     }
 
-    private val sessionFactory = hibernateConfig.sessionFactoryForRegisteredSchemas()
-    private val criteriaBuilder = sessionFactory.criteriaBuilder
+    private var sessionFactory = hibernateConfig.sessionFactoryForRegisteredSchemas()
+    private var criteriaBuilder = sessionFactory.criteriaBuilder
 
     /**
      * Maintain a list of contract state interfaces to concrete types stored in the vault
@@ -64,6 +64,10 @@ class HibernateVaultQueryImpl(hibernateConfig: HibernateConfiguration,
     override fun <T : ContractState> _queryBy(criteria: QueryCriteria, paging: PageSpecification, sorting: Sort, contractType: Class<out T>): Vault.Page<T> {
         log.info("Vault Query for contract type: $contractType, criteria: $criteria, pagination: $paging, sorting: $sorting")
 
+        // refresh to include any schemas registered after initial VQ service initialisation
+        sessionFactory = hibernateConfig.sessionFactoryForRegisteredSchemas()
+        criteriaBuilder = sessionFactory.criteriaBuilder
+
         // calculate total results where a page specification has been defined
         var totalStates = -1L
         if (!paging.isDefault) {

node/src/test/java/net/corda/node/services/vault/VaultQueryJavaTests.java

@@ -65,8 +65,7 @@ public class VaultQueryJavaTests extends TestDependencyInjectionBase {
         ArrayList<KeyPair> keys = new ArrayList<>();
         keys.add(getMEGA_CORP_KEY());
         keys.add(getDUMMY_NOTARY_KEY());
-        Set<MappedSchema> requiredSchemas = new HashSet<>();
-        requiredSchemas.add(CashSchemaV1.INSTANCE);
+        Set<MappedSchema> requiredSchemas = Collections.singleton(CashSchemaV1.INSTANCE);
         IdentityService identitySvc = makeTestIdentityService();
         @SuppressWarnings("unchecked")
         Pair<CordaPersistence, MockServices> databaseAndServices = makeTestDatabaseAndMockServices(requiredSchemas, keys, () -> identitySvc);

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

@@ -43,10 +43,7 @@ class DBTransactionStorageTests : TestDependencyInjectionBase() {
         LogHelper.setLevel(PersistentUniquenessProvider::class)
         val dataSourceProps = makeTestDataSourceProperties()
 
-        val transactionSchema = MappedSchema(schemaFamily = javaClass, version = 1,
-                mappedTypes = listOf(DBTransactionStorage.DBTransaction::class.java))
-
-        val createSchemaService = { NodeSchemaService(setOf(VaultSchemaV1, CashSchemaV1, SampleCashSchemaV2, SampleCashSchemaV3, transactionSchema)) }
+        val createSchemaService = { NodeSchemaService() }
 
         database = configureDatabase(dataSourceProps, makeTestDatabaseProperties(), createSchemaService, ::makeTestIdentityService)
 

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

@@ -76,8 +76,7 @@ class HibernateConfigurationTest : TestDependencyInjectionBase() {
         issuerServices = MockServices(DUMMY_CASH_ISSUER_KEY, BOB_KEY, BOC_KEY)
         val dataSourceProps = makeTestDataSourceProperties()
         val defaultDatabaseProperties = makeTestDatabaseProperties()
-        val customSchemas = setOf(VaultSchemaV1, CashSchemaV1, SampleCashSchemaV2, SampleCashSchemaV3)
-        val createSchemaService = { NodeSchemaService(customSchemas) }
+        val createSchemaService = { NodeSchemaService() }
         database = configureDatabase(dataSourceProps, defaultDatabaseProperties, createSchemaService, ::makeTestIdentityService)
         database.transaction {
             hibernateConfig = database.hibernateConfig
@@ -97,6 +96,7 @@ class HibernateConfigurationTest : TestDependencyInjectionBase() {
         }
         setUpDb()
 
+        val customSchemas = setOf(VaultSchemaV1, CashSchemaV1, SampleCashSchemaV2, SampleCashSchemaV3)
         sessionFactory = hibernateConfig.sessionFactoryForSchemas(*customSchemas.toTypedArray())
         entityManager = sessionFactory.createEntityManager()
         criteriaBuilder = sessionFactory.criteriaBuilder

node/src/test/kotlin/net/corda/node/services/schema/HibernateObserverTests.kt

@@ -7,22 +7,19 @@ import net.corda.core.node.services.Vault
 import net.corda.core.schemas.MappedSchema
 import net.corda.core.schemas.PersistentState
 import net.corda.core.schemas.QueryableState
-import net.corda.testing.LogHelper
 import net.corda.node.services.api.SchemaService
 import net.corda.node.utilities.DatabaseTransactionManager
 import net.corda.node.utilities.configureDatabase
+import net.corda.testing.LogHelper
 import net.corda.testing.MEGA_CORP
 import net.corda.testing.contracts.DUMMY_PROGRAM_ID
 import net.corda.testing.node.MockServices.Companion.makeTestDataSourceProperties
 import net.corda.testing.node.MockServices.Companion.makeTestDatabaseProperties
 import net.corda.testing.node.MockServices.Companion.makeTestIdentityService
-import org.hibernate.annotations.Cascade
-import org.hibernate.annotations.CascadeType
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import rx.subjects.PublishSubject
-import javax.persistence.*
 import kotlin.test.assertEquals
 
 
@@ -38,32 +35,6 @@ class HibernateObserverTests {
         LogHelper.reset(HibernateObserver::class)
     }
 
-    class SchemaFamily
-
-    @Entity
-    @Table(name = "Parents")
-    class Parent : PersistentState() {
-        @OneToMany(fetch = FetchType.LAZY)
-        @JoinColumns(JoinColumn(name = "transaction_id", referencedColumnName = "transaction_id"), JoinColumn(name = "output_index", referencedColumnName = "output_index"))
-        @OrderColumn
-        @Cascade(CascadeType.PERSIST)
-        var children: MutableSet<Child> = mutableSetOf()
-    }
-
-    @Suppress("unused")
-    @Entity
-    @Table(name = "Children")
-    class Child {
-        @Id
-        @GeneratedValue(strategy = GenerationType.IDENTITY)
-        @Column(name = "child_id", unique = true, nullable = false)
-        var childId: Int? = null
-
-        @ManyToOne(fetch = FetchType.LAZY)
-        @JoinColumns(JoinColumn(name = "transaction_id", referencedColumnName = "transaction_id"), JoinColumn(name = "output_index", referencedColumnName = "output_index"))
-        var parent: Parent? = null
-    }
-
     class TestState : QueryableState {
         override fun supportedSchemas(): Iterable<MappedSchema> {
             throw UnsupportedOperationException()
@@ -80,17 +51,19 @@ class HibernateObserverTests {
     // This method does not use back quotes for a nice name since it seems to kill the kotlin compiler.
     @Test
     fun testChildObjectsArePersisted() {
-        val testSchema = object : MappedSchema(SchemaFamily::class.java, 1, setOf(Parent::class.java, Child::class.java)) {}
+        val testSchema = TestSchema
         val rawUpdatesPublisher = PublishSubject.create<Vault.Update<ContractState>>()
         val schemaService = object : SchemaService {
+            override fun registerCustomSchemas(customSchemas: Set<MappedSchema>) {}
+
             override val schemaOptions: Map<MappedSchema, SchemaService.SchemaOptions> = emptyMap()
 
             override fun selectSchemas(state: ContractState): Iterable<MappedSchema> = setOf(testSchema)
 
             override fun generateMappedObject(state: ContractState, schema: MappedSchema): PersistentState {
-                val parent = Parent()
-                parent.children.add(Child())
-                parent.children.add(Child())
+                val parent = TestSchema.Parent()
+                parent.children.add(TestSchema.Child())
+                parent.children.add(TestSchema.Child())
                 return parent
             }
         }

node/src/test/kotlin/net/corda/node/services/schema/NodeSchemaServiceTest.kt

@@ -0,0 +1,89 @@
+package net.corda.node.services.schema
+
+import co.paralleluniverse.fibers.Suspendable
+import net.corda.core.flows.FlowLogic
+import net.corda.core.flows.StartableByRPC
+import net.corda.core.messaging.startFlow
+import net.corda.core.schemas.MappedSchema
+import net.corda.core.schemas.PersistentState
+import net.corda.core.utilities.getOrThrow
+import net.corda.node.services.api.ServiceHubInternal
+import net.corda.testing.driver.driver
+import net.corda.testing.node.MockNetwork
+import net.corda.testing.schemas.DummyLinearStateSchemaV1
+import org.hibernate.annotations.Cascade
+import org.hibernate.annotations.CascadeType
+import org.junit.Test
+import javax.persistence.*
+import kotlin.test.assertTrue
+
+class NodeSchemaServiceTest {
+    /**
+     * Note: this test requires explicitly registering custom contract schemas with a MockNode
+     */
+    @Test
+    fun `registering custom schemas for testing with MockNode`() {
+        val mockNet = MockNetwork()
+        val mockNode = mockNet.createNode()
+        mockNet.runNetwork()
+
+        mockNode.registerCustomSchemas(setOf(DummyLinearStateSchemaV1))
+        val schemaService = mockNode.services.schemaService
+        assertTrue(schemaService.schemaOptions.containsKey(DummyLinearStateSchemaV1))
+
+        mockNet.stopNodes()
+    }
+
+    /**
+     * Note: this test verifies auto-scanning to register identified [MappedSchema] schemas.
+     *       By default, Driver uses the caller package for auto-scanning:
+     *       System.setProperty("net.corda.node.cordapp.scan.package", callerPackage)
+     */
+    @Test
+    fun `auto scanning of custom schemas for testing with Driver`() {
+        driver (startNodesInProcess = true) {
+            val node = startNode()
+            val nodeHandle = node.getOrThrow()
+            val result = nodeHandle.rpc.startFlow(::MappedSchemasFlow)
+            val mappedSchemas = result.returnValue.getOrThrow()
+            assertTrue(mappedSchemas.contains(TestSchema.name))
+        }
+    }
+
+    @StartableByRPC
+    class MappedSchemasFlow : FlowLogic<List<String>>() {
+        @Suspendable
+        override fun call() : List<String> {
+            // returning MappedSchema's as String'ified family names to avoid whitelist serialization errors
+            return (this.serviceHub as ServiceHubInternal).schemaService.schemaOptions.keys.map { it.name }
+        }
+    }
+}
+
+class SchemaFamily
+
+object TestSchema : MappedSchema(SchemaFamily::class.java, 1, setOf(Parent::class.java, Child::class.java)) {
+    @Entity
+    @Table(name = "Parents")
+    class Parent : PersistentState() {
+        @OneToMany(fetch = FetchType.LAZY)
+        @JoinColumns(JoinColumn(name = "transaction_id", referencedColumnName = "transaction_id"), JoinColumn(name = "output_index", referencedColumnName = "output_index"))
+        @OrderColumn
+        @Cascade(CascadeType.PERSIST)
+        var children: MutableSet<Child> = mutableSetOf()
+    }
+
+    @Suppress("unused")
+    @Entity
+    @Table(name = "Children")
+    class Child {
+        @Id
+        @GeneratedValue(strategy = GenerationType.IDENTITY)
+        @Column(name = "child_id", unique = true, nullable = false)
+        var childId: Int? = null
+
+        @ManyToOne(fetch = FetchType.LAZY)
+        @JoinColumns(JoinColumn(name = "transaction_id", referencedColumnName = "transaction_id"), JoinColumn(name = "output_index", referencedColumnName = "output_index"))
+        var parent: Parent? = null
+    }
+}

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

@@ -23,6 +23,7 @@ import net.corda.finance.contracts.asset.Cash
 import net.corda.finance.contracts.asset.DUMMY_CASH_ISSUER
 import net.corda.finance.contracts.asset.DUMMY_CASH_ISSUER_KEY
 import net.corda.finance.contracts.getCashBalance
+import net.corda.finance.schemas.CashSchemaV1
 import net.corda.finance.utils.sumCash
 import net.corda.node.utilities.CordaPersistence
 import net.corda.testing.*
@@ -53,7 +54,8 @@ class NodeVaultServiceTest : TestDependencyInjectionBase() {
     @Before
     fun setUp() {
         LogHelper.setLevel(NodeVaultService::class)
-        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(BOC_KEY, DUMMY_CASH_ISSUER_KEY))
+        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(BOC_KEY, DUMMY_CASH_ISSUER_KEY),
+                                                                  customSchemas = setOf(CashSchemaV1))
         database = databaseAndServices.first
         services = databaseAndServices.second
         issuerServices = MockServices(DUMMY_CASH_ISSUER_KEY, BOC_KEY)

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

@@ -63,7 +63,9 @@ class VaultQueryTests : TestDependencyInjectionBase() {
         // register additional identities
         identitySvc.verifyAndRegisterIdentity(CASH_NOTARY_IDENTITY)
         identitySvc.verifyAndRegisterIdentity(BOC_IDENTITY)
-        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(MEGA_CORP_KEY, DUMMY_NOTARY_KEY), createIdentityService = { identitySvc })
+        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(MEGA_CORP_KEY, DUMMY_NOTARY_KEY),
+                                                                  createIdentityService = { identitySvc },
+                                                                  customSchemas = setOf(CashSchemaV1, CommercialPaperSchemaV1, DummyLinearStateSchemaV1))
         database = databaseAndServices.first
         services = databaseAndServices.second
         notaryServices = MockServices(DUMMY_NOTARY_KEY, DUMMY_CASH_ISSUER_KEY, BOC_KEY, MEGA_CORP_KEY)

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

@@ -16,6 +16,7 @@ import net.corda.finance.contracts.asset.Cash
 import net.corda.finance.contracts.asset.DUMMY_CASH_ISSUER
 import net.corda.finance.contracts.asset.DUMMY_CASH_ISSUER_KEY
 import net.corda.finance.contracts.getCashBalance
+import net.corda.finance.schemas.CashSchemaV1
 import net.corda.node.utilities.CordaPersistence
 import net.corda.testing.*
 import net.corda.testing.contracts.*
@@ -44,7 +45,8 @@ class VaultWithCashTest : TestDependencyInjectionBase() {
     @Before
     fun setUp() {
         LogHelper.setLevel(VaultWithCashTest::class)
-        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(DUMMY_CASH_ISSUER_KEY, DUMMY_NOTARY_KEY))
+        val databaseAndServices = makeTestDatabaseAndMockServices(keys = listOf(DUMMY_CASH_ISSUER_KEY, DUMMY_NOTARY_KEY),
+                                                                  customSchemas = setOf(CashSchemaV1))
         database = databaseAndServices.first
         services = databaseAndServices.second
         issuerServices = MockServices(DUMMY_CASH_ISSUER_KEY, MEGA_CORP_KEY)

samples/trader-demo/src/integration-test/kotlin/net/corda/traderdemo/TraderDemoTest.kt

@@ -7,6 +7,8 @@ import net.corda.core.utilities.millis
 import net.corda.finance.DOLLARS
 import net.corda.finance.flows.CashIssueFlow
 import net.corda.finance.flows.CashPaymentFlow
+import net.corda.finance.schemas.CashSchemaV1
+import net.corda.finance.schemas.CommercialPaperSchemaV1
 import net.corda.node.services.FlowPermissions.Companion.startFlowPermission
 import net.corda.node.services.transactions.SimpleNotaryService
 import net.corda.nodeapi.User
@@ -38,6 +40,8 @@ class TraderDemoTest : NodeBasedTest() {
         val (nodeA, nodeB, bankNode) = listOf(nodeAFuture, nodeBFuture, bankNodeFuture, notaryFuture).map { it.getOrThrow() }
 
         nodeA.registerInitiatedFlow(BuyerFlow::class.java)
+        nodeA.registerCustomSchemas(setOf(CashSchemaV1))
+        nodeB.registerCustomSchemas(setOf(CashSchemaV1, CommercialPaperSchemaV1))
 
         val (nodeARpc, nodeBRpc) = listOf(nodeA, nodeB).map {
             val client = CordaRPCClient(it.configuration.rpcAddress!!, initialiseSerialization = false)

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

@@ -12,8 +12,6 @@ import net.corda.core.serialization.SerializeAsToken
 import net.corda.core.serialization.SingletonSerializeAsToken
 import net.corda.core.transactions.SignedTransaction
 import net.corda.core.utilities.NonEmptySet
-import net.corda.finance.schemas.CashSchemaV1
-import net.corda.finance.schemas.CommercialPaperSchemaV1
 import net.corda.node.VersionInfo
 import net.corda.node.services.api.StateMachineRecordedTransactionMappingStorage
 import net.corda.node.services.api.WritableTransactionStorage
@@ -25,13 +23,11 @@ import net.corda.node.services.persistence.InMemoryStateMachineRecordedTransacti
 import net.corda.node.services.schema.HibernateObserver
 import net.corda.node.services.schema.NodeSchemaService
 import net.corda.node.services.transactions.InMemoryTransactionVerifierService
-import net.corda.node.services.upgrade.ContractUpgradeServiceImpl
 import net.corda.node.services.vault.HibernateVaultQueryImpl
 import net.corda.node.services.vault.NodeVaultService
 import net.corda.node.utilities.CordaPersistence
 import net.corda.node.utilities.configureDatabase
 import net.corda.testing.*
-import net.corda.testing.schemas.DummyLinearStateSchemaV1
 import org.bouncycastle.operator.ContentSigner
 import rx.Observable
 import rx.subjects.PublishSubject
@@ -102,7 +98,7 @@ open class MockServices(vararg val keys: KeyPair) : ServiceHub {
          * @return a pair where the first element is the instance of [CordaPersistence] and the second is [MockServices].
          */
         @JvmStatic
-        fun makeTestDatabaseAndMockServices(customSchemas: Set<MappedSchema> = setOf(CommercialPaperSchemaV1, DummyLinearStateSchemaV1, CashSchemaV1),
+        fun makeTestDatabaseAndMockServices(customSchemas: Set<MappedSchema> = emptySet(),
                                             keys: List<KeyPair> = listOf(MEGA_CORP_KEY),
                                             createIdentityService: () -> IdentityService = { makeTestIdentityService() }): Pair<CordaPersistence, MockServices> {
             val dataSourceProps = makeTestDataSourceProperties()