merged OS into master

testing/test-utils/src/main/kotlin/net/corda/testing/internal/MockCordappProvider.kt

@@ -2,6 +2,7 @@ package net.corda.testing.internal
 
 import net.corda.core.contracts.ContractClassName
 import net.corda.core.cordapp.Cordapp
+import net.corda.core.internal.TEST_UPLOADER
 import net.corda.core.internal.cordapp.CordappImpl
 import net.corda.core.node.services.AttachmentId
 import net.corda.core.node.services.AttachmentStorage
@@ -14,9 +15,10 @@ import java.util.*
 class MockCordappProvider(
         cordappLoader: CordappLoader,
         attachmentStorage: AttachmentStorage,
-        val cordappConfigProvider: MockCordappConfigProvider = MockCordappConfigProvider()
-) : CordappProviderImpl(cordappLoader, cordappConfigProvider, attachmentStorage) {
-    constructor(cordappLoader: CordappLoader, attachmentStorage: AttachmentStorage) : this(cordappLoader, attachmentStorage, MockCordappConfigProvider())
+        whitelistedContractImplementations: Map<String, List<AttachmentId>>,
+        cordappConfigProvider: MockCordappConfigProvider = MockCordappConfigProvider()
+) : CordappProviderImpl(cordappLoader, cordappConfigProvider, attachmentStorage, whitelistedContractImplementations) {
+    constructor(cordappLoader: CordappLoader, attachmentStorage: AttachmentStorage, whitelistedContractImplementations: Map<String, List<AttachmentId>>) : this(cordappLoader, attachmentStorage, whitelistedContractImplementations, MockCordappConfigProvider())
 
     val cordappRegistry = mutableListOf<Pair<Cordapp, AttachmentId>>()
 
@@ -33,20 +35,21 @@ class MockCordappProvider(
                 customSchemas = emptySet(),
                 jarPath = Paths.get("").toUri().toURL())
         if (cordappRegistry.none { it.first.contractClassNames.contains(contractClassName) }) {
-            cordappRegistry.add(Pair(cordapp, findOrImportAttachment(contractClassName.toByteArray(), attachments)))
+            cordappRegistry.add(Pair(cordapp, findOrImportAttachment(listOf(contractClassName), contractClassName.toByteArray(), attachments)))
         }
     }
 
-    override fun getContractAttachmentID(contractClassName: ContractClassName): AttachmentId? = cordappRegistry.find { it.first.contractClassNames.contains(contractClassName) }?.second ?: super.getContractAttachmentID(contractClassName)
+    override fun getContractAttachmentID(contractClassName: ContractClassName): AttachmentId? = cordappRegistry.find { it.first.contractClassNames.contains(contractClassName) }?.second
+            ?: super.getContractAttachmentID(contractClassName)
 
-    private fun findOrImportAttachment(data: ByteArray, attachments: MockAttachmentStorage): AttachmentId {
+    private fun findOrImportAttachment(contractClassNames: List<ContractClassName>, data: ByteArray, attachments: MockAttachmentStorage): AttachmentId {
         val existingAttachment = attachments.files.filter {
-            Arrays.equals(it.value, data)
+            Arrays.equals(it.value.second, data)
         }
         return if (!existingAttachment.isEmpty()) {
             existingAttachment.keys.first()
         } else {
-            attachments.importAttachment(data.inputStream())
+            attachments.importContractAttachment(contractClassNames, TEST_UPLOADER, data.inputStream())
         }
     }
 }

testing/test-utils/src/main/kotlin/net/corda/testing/services/MockAttachmentStorage.kt

@@ -1,14 +1,18 @@
 package net.corda.testing.services
 
 import net.corda.core.contracts.Attachment
+import net.corda.core.contracts.ContractAttachment
+import net.corda.core.contracts.ContractClassName
 import net.corda.core.crypto.SecureHash
 import net.corda.core.crypto.sha256
 import net.corda.core.internal.AbstractAttachment
+import net.corda.core.internal.UNKNOWN_UPLOADER
 import net.corda.core.node.services.AttachmentId
 import net.corda.core.node.services.AttachmentStorage
 import net.corda.core.node.services.vault.AttachmentQueryCriteria
 import net.corda.core.node.services.vault.AttachmentSort
 import net.corda.core.serialization.SingletonSerializeAsToken
+import net.corda.nodeapi.internal.withContractsInJar
 import java.io.ByteArrayOutputStream
 import java.io.InputStream
 import java.util.*
@@ -24,31 +28,17 @@ class MockAttachmentStorage : AttachmentStorage, SingletonSerializeAsToken() {
         }
     }
 
-    override fun importAttachment(jar: InputStream): AttachmentId {
-        // JIS makes read()/readBytes() return bytes of the current file, but we want to hash the entire container here.
-        require(jar !is JarInputStream)
+    val files = HashMap<SecureHash, Pair<Attachment, ByteArray>>()
 
-        val bytes = getBytes(jar)
+    override fun importAttachment(jar: InputStream): AttachmentId = importAttachment(jar, UNKNOWN_UPLOADER, null)
 
-        val sha256 = bytes.sha256()
-        if (!files.containsKey(sha256)) {
-            files[sha256] = bytes
+    override fun importAttachment(jar: InputStream, uploader: String, filename: String?): AttachmentId {
+        return withContractsInJar(jar) { contractClassNames, inputStream ->
+            importAttachmentInternal(inputStream, uploader, filename, contractClassNames)
         }
-        return sha256
     }
 
-    override fun importAttachment(jar: InputStream, uploader: String, filename: String): AttachmentId {
-        return importAttachment(jar)
-    }
-
-    val files = HashMap<SecureHash, ByteArray>()
-
-    private class MockAttachment(dataLoader: () -> ByteArray, override val id: SecureHash) : AbstractAttachment(dataLoader)
-
-    override fun openAttachment(id: SecureHash): Attachment? {
-        val f = files[id] ?: return null
-        return MockAttachment({ f }, id)
-    }
+    override fun openAttachment(id: SecureHash): Attachment? = files[id]?.first
 
     override fun queryAttachments(criteria: AttachmentQueryCriteria, sorting: AttachmentSort?): List<AttachmentId> {
         throw NotImplementedError("Querying for attachments not implemented")
@@ -56,11 +46,6 @@ class MockAttachmentStorage : AttachmentStorage, SingletonSerializeAsToken() {
 
     override fun hasAttachment(attachmentId: AttachmentId) = files.containsKey(attachmentId)
 
-    fun getAttachmentIdAndBytes(jar: InputStream): Pair<AttachmentId, ByteArray> {
-        val bytes = getBytes(jar)
-        return Pair(bytes.sha256(), bytes)
-    }
-
     override fun importOrGetAttachment(jar: InputStream): AttachmentId {
         try {
             return importAttachment(jar)
@@ -68,4 +53,25 @@ class MockAttachmentStorage : AttachmentStorage, SingletonSerializeAsToken() {
             return AttachmentId.parse(faee.message!!)
         }
     }
+
+    fun importContractAttachment(contractClassNames: List<ContractClassName>, uploader: String, jar: InputStream): AttachmentId = importAttachmentInternal(jar, uploader, null, contractClassNames)
+
+    fun getAttachmentIdAndBytes(jar: InputStream): Pair<AttachmentId, ByteArray> = getBytes(jar).let { bytes -> Pair(bytes.sha256(), bytes) }
+
+    private class MockAttachment(dataLoader: () -> ByteArray, override val id: SecureHash) : AbstractAttachment(dataLoader)
+
+    private fun importAttachmentInternal(jar: InputStream, uploader: String, filename: String?, contractClassNames: List<ContractClassName>? = null): AttachmentId {
+        // JIS makes read()/readBytes() return bytes of the current file, but we want to hash the entire container here.
+        require(jar !is JarInputStream)
+
+        val bytes = getBytes(jar)
+
+        val sha256 = bytes.sha256()
+        if (sha256 !in files.keys) {
+            val baseAttachment = MockAttachment({ bytes }, sha256)
+            val attachment = if (contractClassNames == null || contractClassNames.isEmpty()) baseAttachment else ContractAttachment(baseAttachment, contractClassNames.first(), contractClassNames.toSet(), uploader)
+            files[sha256] = Pair(attachment, bytes)
+        }
+        return sha256
+    }
 }