From 278d5b05e6475097710cb96c44749e75ec4de1a7 Mon Sep 17 00:00:00 2001
From: Christian Basler <chrigu.meyer@gmail.com>
Date: Sun, 26 Nov 2017 20:30:05 +0100
Subject: [PATCH] Split LabelRepository off the MessageRepository

---
 .../ch/dissem/bitmessage/BitmessageContext.kt |  11 ++
 .../ch/dissem/bitmessage/InternalContext.kt   |  11 +-
 .../ch/dissem/bitmessage/entity/Plaintext.kt  | 133 ++++++++------
 .../ports/AbstractLabelRepository.kt          |  33 ++++
 .../ports/AbstractMessageRepository.kt        |  22 +--
 .../dissem/bitmessage/ports/DefaultLabeler.kt |  14 +-
 .../bitmessage/ports/LabelRepository.kt       |  27 +++
 .../bitmessage/ports/MessageRepository.kt     |   6 -
 .../ch/dissem/bitmessage/utils/Collections.kt |   9 +
 .../bitmessage/BitmessageContextTest.kt       |   1 +
 .../ch/dissem/bitmessage/utils/TestUtils.kt   |   2 +
 .../dissem/bitmessage/demo/Application.java   |   2 +-
 .../java/ch/dissem/bitmessage/SystemTest.java |   6 +-
 .../networking/NetworkHandlerTest.kt          |   2 +
 .../repository/JdbcLabelRepository.kt         | 123 +++++++++++++
 .../repository/JdbcMessageRepository.kt       | 166 +++++-------------
 .../repository/JdbcLabelRepositoryTest.kt     |  49 ++++++
 .../repository/JdbcMessageRepositoryTest.kt   |  24 +--
 .../dissem/bitmessage/wif/WifExporterTest.kt  |   1 +
 .../dissem/bitmessage/wif/WifImporterTest.kt  |   1 +
 20 files changed, 407 insertions(+), 236 deletions(-)
 create mode 100644 core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractLabelRepository.kt
 create mode 100644 core/src/main/kotlin/ch/dissem/bitmessage/ports/LabelRepository.kt
 create mode 100644 repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepository.kt
 create mode 100644 repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepositoryTest.kt

diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/BitmessageContext.kt b/core/src/main/kotlin/ch/dissem/bitmessage/BitmessageContext.kt
index 94a29ee..a97cfdc 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/BitmessageContext.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/BitmessageContext.kt
@@ -73,6 +73,9 @@ class BitmessageContext private constructor(builder: BitmessageContext.Builder)
     val addresses: AddressRepository
         @JvmName("addresses") get
 
+    val labels: LabelRepository
+        @JvmName("labels") get
+
     val messages: MessageRepository
         @JvmName("messages") get
 
@@ -301,6 +304,7 @@ class BitmessageContext private constructor(builder: BitmessageContext.Builder)
         var nodeRegistry by Delegates.notNull<NodeRegistry>()
         var networkHandler by Delegates.notNull<NetworkHandler>()
         var addressRepo by Delegates.notNull<AddressRepository>()
+        var labelRepo by Delegates.notNull<LabelRepository>()
         var messageRepo by Delegates.notNull<MessageRepository>()
         var proofOfWorkRepo by Delegates.notNull<ProofOfWorkRepository>()
         var proofOfWorkEngine: ProofOfWorkEngine? = null
@@ -330,6 +334,11 @@ class BitmessageContext private constructor(builder: BitmessageContext.Builder)
             return this
         }
 
+        fun labelRepo(labelRepo: LabelRepository): Builder {
+            this.labelRepo = labelRepo
+            return this
+        }
+
         fun messageRepo(messageRepo: MessageRepository): Builder {
             this.messageRepo = messageRepo
             return this
@@ -386,6 +395,7 @@ class BitmessageContext private constructor(builder: BitmessageContext.Builder)
             builder.nodeRegistry,
             builder.networkHandler,
             builder.addressRepo,
+            builder.labelRepo,
             builder.messageRepo,
             builder.proofOfWorkRepo,
             builder.proofOfWorkEngine ?: MultiThreadedPOWEngine(),
@@ -400,6 +410,7 @@ class BitmessageContext private constructor(builder: BitmessageContext.Builder)
             builder.preferences
         )
         this.addresses = builder.addressRepo
+        this.labels = builder.labelRepo
         this.messages = builder.messageRepo
         (builder.listener as? Listener.WithContext)?.setContext(this)
         internals.proofOfWorkService.doMissingProofOfWork(builder.preferences.doMissingProofOfWorkDelayInSeconds * 1000L)
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/InternalContext.kt b/core/src/main/kotlin/ch/dissem/bitmessage/InternalContext.kt
index 6415b8d..c7b303d 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/InternalContext.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/InternalContext.kt
@@ -40,11 +40,12 @@ import java.util.concurrent.Executors
  */
 class InternalContext(
     val cryptography: Cryptography,
-    val inventory: ch.dissem.bitmessage.ports.Inventory,
+    val inventory: Inventory,
     val nodeRegistry: NodeRegistry,
     val networkHandler: NetworkHandler,
     val addressRepository: AddressRepository,
-    val messageRepository: ch.dissem.bitmessage.ports.MessageRepository,
+    val labelRepository: LabelRepository,
+    val messageRepository: MessageRepository,
     val proofOfWorkRepository: ProofOfWorkRepository,
     val proofOfWorkEngine: ProofOfWorkEngine,
     val customCommandHandler: CustomCommandHandler,
@@ -216,7 +217,9 @@ class InternalContext(
     companion object {
         private val LOG = LoggerFactory.getLogger(InternalContext::class.java)
 
-        @JvmField val NETWORK_NONCE_TRIALS_PER_BYTE: Long = 1000
-        @JvmField val NETWORK_EXTRA_BYTES: Long = 1000
+        @JvmField
+        val NETWORK_NONCE_TRIALS_PER_BYTE: Long = 1000
+        @JvmField
+        val NETWORK_EXTRA_BYTES: Long = 1000
     }
 }
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/entity/Plaintext.kt b/core/src/main/kotlin/ch/dissem/bitmessage/entity/Plaintext.kt
index 236d581..d382436 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/entity/Plaintext.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/entity/Plaintext.kt
@@ -35,6 +35,7 @@ import java.nio.ByteBuffer
 import java.util.*
 import java.util.Collections
 import kotlin.collections.HashSet
+import kotlin.collections.LinkedHashSet
 
 private fun message(encoding: Plaintext.Encoding, subject: String, body: String): ByteArray = when (encoding) {
     SIMPLE -> "Subject:$subject\nBody:$body".toByteArray()
@@ -254,7 +255,7 @@ class Plaintext private constructor(
         received = builder.received,
         initialHash = null,
         ttl = builder.ttl,
-        labels = builder.labels,
+        labels = LinkedHashSet(builder.labels),
         status = builder.status ?: Status.RECEIVED
     ) {
         id = builder.id
@@ -278,28 +279,30 @@ class Plaintext private constructor(
         get() {
             val s = Scanner(ByteArrayInputStream(message), "UTF-8")
             val firstLine = s.nextLine()
-            if (encodingCode == EXTENDED.code) {
-                if (Message.TYPE == extendedData?.type) {
-                    return (extendedData!!.content as? Message)?.subject
+            return when (encodingCode) {
+                EXTENDED.code -> if (Message.TYPE == extendedData?.type) {
+                    (extendedData!!.content as? Message)?.subject
                 } else {
-                    return null
+                    null
+                }
+                SIMPLE.code -> firstLine.substring("Subject:".length).trim { it <= ' ' }
+                else -> {
+                    if (firstLine.length > 50) {
+                        firstLine.substring(0, 50).trim { it <= ' ' } + "..."
+                    } else {
+                        firstLine
+                    }
                 }
-            } else if (encodingCode == SIMPLE.code) {
-                return firstLine.substring("Subject:".length).trim { it <= ' ' }
-            } else if (firstLine.length > 50) {
-                return firstLine.substring(0, 50).trim { it <= ' ' } + "..."
-            } else {
-                return firstLine
             }
         }
 
     val text: String?
         get() {
             if (encodingCode == EXTENDED.code) {
-                if (Message.TYPE == extendedData?.type) {
-                    return (extendedData?.content as Message?)?.body
+                return if (Message.TYPE == extendedData?.type) {
+                    (extendedData?.content as Message?)?.body
                 } else {
-                    return null
+                    null
                 }
             } else {
                 val text = String(message)
@@ -322,20 +325,20 @@ class Plaintext private constructor(
     val parents: List<InventoryVector>
         get() {
             val extendedData = extendedData ?: return emptyList()
-            if (Message.TYPE == extendedData.type) {
-                return (extendedData.content as Message).parents
+            return if (Message.TYPE == extendedData.type) {
+                (extendedData.content as Message).parents
             } else {
-                return emptyList()
+                emptyList()
             }
         }
 
     val files: List<Attachment>
         get() {
             val extendedData = extendedData ?: return emptyList()
-            if (Message.TYPE == extendedData.type) {
-                return (extendedData.content as Message).files
+            return if (Message.TYPE == extendedData.type) {
+                (extendedData.content as Message).files
             } else {
-                return emptyList()
+                emptyList()
             }
         }
 
@@ -378,8 +381,8 @@ class Plaintext private constructor(
 
     override fun toString(): String {
         val subject = subject
-        if (subject?.isNotEmpty() ?: false) {
-            return subject!!
+        if (subject?.isNotEmpty() == true) {
+            return subject
         } else {
             return Strings.hex(
                 initialHash ?: return super.toString()
@@ -529,32 +532,43 @@ class Plaintext private constructor(
     }
 
     class Builder(internal val type: Type) {
-        internal var id: Any? = null
-        internal var inventoryVector: InventoryVector? = null
-        internal var from: BitmessageAddress? = null
-        internal var to: BitmessageAddress? = null
-        private var addressVersion: Long = 0
-        private var stream: Long = 0
-        private var behaviorBitfield: Int = 0
-        private var publicSigningKey: ByteArray? = null
-        private var publicEncryptionKey: ByteArray? = null
-        private var nonceTrialsPerByte: Long = 0
-        private var extraBytes: Long = 0
-        private var destinationRipe: ByteArray? = null
-        private var preventAck: Boolean = false
-        internal var encoding: Long = 0
-        internal var message = ByteArray(0)
-        internal var ackData: ByteArray? = null
-        internal var ackMessage: ByteArray? = null
-        internal var signature: ByteArray? = null
-        internal var sent: Long? = null
-        internal var received: Long? = null
-        internal var status: Status? = null
-        internal val labels = LinkedHashSet<Label>()
-        internal var ttl: Long = 0
-        internal var retries: Int = 0
-        internal var nextTry: Long? = null
-        internal var conversation: UUID? = null
+        var id: Any? = null
+        var inventoryVector: InventoryVector? = null
+        var from: BitmessageAddress? = null
+        var to: BitmessageAddress? = null
+            set(value) {
+                if (value != null) {
+                    if (type != MSG && to != null)
+                        throw IllegalArgumentException("recipient address only allowed for msg")
+                    field = value
+                }
+            }
+        var addressVersion: Long = 0
+        var stream: Long = 0
+        var behaviorBitfield: Int = 0
+        var publicSigningKey: ByteArray? = null
+        var publicEncryptionKey: ByteArray? = null
+        var nonceTrialsPerByte: Long = 0
+        var extraBytes: Long = 0
+        var destinationRipe: ByteArray? = null
+            set(value) {
+                if (type != MSG && value != null) throw IllegalArgumentException("ripe only allowed for msg")
+                field = value
+            }
+        var preventAck: Boolean = false
+        var encoding: Long = 0
+        var message = ByteArray(0)
+        var ackData: ByteArray? = null
+        var ackMessage: ByteArray? = null
+        var signature: ByteArray? = null
+        var sent: Long? = null
+        var received: Long? = null
+        var status: Status? = null
+        var labels: Collection<Label> = emptySet()
+        var ttl: Long = 0
+        var retries: Int = 0
+        var nextTry: Long? = null
+        var conversation: UUID? = null
 
         fun id(id: Any): Builder {
             this.id = id
@@ -572,11 +586,7 @@ class Plaintext private constructor(
         }
 
         fun to(address: BitmessageAddress?): Builder {
-            if (address != null) {
-                if (type != MSG && to != null)
-                    throw IllegalArgumentException("recipient address only allowed for msg")
-                to = address
-            }
+            to = address
             return this
         }
 
@@ -616,7 +626,6 @@ class Plaintext private constructor(
         }
 
         fun destinationRipe(ripe: ByteArray?): Builder {
-            if (type != MSG && ripe != null) throw IllegalArgumentException("ripe only allowed for msg")
             this.destinationRipe = ripe
             return this
         }
@@ -692,7 +701,7 @@ class Plaintext private constructor(
         }
 
         fun labels(labels: Collection<Label>): Builder {
-            this.labels.addAll(labels)
+            this.labels = labels
             return this
         }
 
@@ -743,6 +752,12 @@ class Plaintext private constructor(
             return this
         }
 
+        @JvmSynthetic
+        inline fun build(block: Builder.() -> Unit): Plaintext {
+            block(this)
+            return build()
+        }
+
         fun build(): Plaintext {
             return Plaintext(this)
         }
@@ -774,5 +789,13 @@ class Plaintext private constructor(
                 .message(Decode.varBytes(input))
                 .ackMessage(if (type == MSG) Decode.varBytes(input) else null)
         }
+
+        @JvmSynthetic
+        inline fun build(type: Type, block: Builder.() -> Unit): Plaintext {
+            val builder = Builder(type)
+            block(builder)
+            return builder.build()
+        }
+
     }
 }
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractLabelRepository.kt b/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractLabelRepository.kt
new file mode 100644
index 0000000..47c7e23
--- /dev/null
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractLabelRepository.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright 2017 Christian Basler
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ch.dissem.bitmessage.ports
+
+import ch.dissem.bitmessage.entity.valueobject.Label
+import ch.dissem.bitmessage.utils.SqlStrings.join
+
+abstract class AbstractLabelRepository : LabelRepository {
+
+    override fun getLabels(): List<Label> {
+        return find("1=1")
+    }
+
+    override fun getLabels(vararg types: Label.Type): List<Label> {
+        return find("type IN (${join(*types)})")
+    }
+
+    protected abstract fun find(where: String): List<Label>
+}
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractMessageRepository.kt b/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractMessageRepository.kt
index 495180d..9df98a9 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractMessageRepository.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/ports/AbstractMessageRepository.kt
@@ -21,8 +21,7 @@ import ch.dissem.bitmessage.entity.BitmessageAddress
 import ch.dissem.bitmessage.entity.Plaintext
 import ch.dissem.bitmessage.entity.valueobject.InventoryVector
 import ch.dissem.bitmessage.entity.valueobject.Label
-import ch.dissem.bitmessage.exception.ApplicationException
-import ch.dissem.bitmessage.utils.SqlStrings.join
+import ch.dissem.bitmessage.utils.Collections.single
 import ch.dissem.bitmessage.utils.Strings
 import ch.dissem.bitmessage.utils.UnixTime
 import java.util.*
@@ -114,25 +113,6 @@ abstract class AbstractMessageRepository : MessageRepository, InternalContext.Co
         return find("conversation=X'${conversationId.toString().replace("-", "")}'")
     }
 
-    override fun getLabels(): List<Label> {
-        return findLabels("1=1")
-    }
-
-    override fun getLabels(vararg types: Label.Type): List<Label> {
-        return findLabels("type IN (${join(*types)})")
-    }
-
-    protected abstract fun findLabels(where: String): List<Label>
-
-
-    protected fun <T> single(collection: Collection<T>): T? {
-        return when (collection.size) {
-            0 -> null
-            1 -> collection.iterator().next()
-            else -> throw ApplicationException("This shouldn't happen, found ${collection.size} items, one or none was expected")
-        }
-    }
-
     /**
      * Finds messages that mach the given where statement, with optional offset and limit. If the limit is set to 0,
      * offset and limit are ignored.
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/ports/DefaultLabeler.kt b/core/src/main/kotlin/ch/dissem/bitmessage/ports/DefaultLabeler.kt
index 01d54d6..4ed8e00 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/ports/DefaultLabeler.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/ports/DefaultLabeler.kt
@@ -35,9 +35,9 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
         msg.status = RECEIVED
         val labelsToAdd =
             if (msg.type == BROADCAST) {
-                ctx.messageRepository.getLabels(Label.Type.BROADCAST, Label.Type.UNREAD)
+                ctx.labelRepository.getLabels(Label.Type.BROADCAST, Label.Type.UNREAD)
             } else {
-                ctx.messageRepository.getLabels(Label.Type.INBOX, Label.Type.UNREAD)
+                ctx.labelRepository.getLabels(Label.Type.INBOX, Label.Type.UNREAD)
             }
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, emptyList())
@@ -45,7 +45,7 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
 
     override fun markAsDraft(msg: Plaintext) {
         msg.status = DRAFT
-        val labelsToAdd = ctx.messageRepository.getLabels(Label.Type.DRAFT)
+        val labelsToAdd = ctx.labelRepository.getLabels(Label.Type.DRAFT)
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, emptyList())
     }
@@ -58,7 +58,7 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
         }
         val labelsToRemove = msg.labels.filter { it.type == Label.Type.DRAFT }
         msg.removeLabel(Label.Type.DRAFT)
-        val labelsToAdd = ctx.messageRepository.getLabels(Label.Type.OUTBOX)
+        val labelsToAdd = ctx.labelRepository.getLabels(Label.Type.OUTBOX)
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, labelsToRemove)
     }
@@ -67,7 +67,7 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
         msg.status = SENT
         val labelsToRemove = msg.labels.filter { it.type == Label.Type.OUTBOX }
         msg.removeLabel(Label.Type.OUTBOX)
-        val labelsToAdd = ctx.messageRepository.getLabels(Label.Type.SENT)
+        val labelsToAdd = ctx.labelRepository.getLabels(Label.Type.SENT)
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, labelsToRemove)
     }
@@ -83,7 +83,7 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
     }
 
     override fun markAsUnread(msg: Plaintext) {
-        val labelsToAdd = ctx.messageRepository.getLabels(Label.Type.UNREAD)
+        val labelsToAdd = ctx.labelRepository.getLabels(Label.Type.UNREAD)
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, emptyList())
     }
@@ -91,7 +91,7 @@ open class DefaultLabeler : Labeler, InternalContext.ContextHolder {
     override fun delete(msg: Plaintext) {
         val labelsToRemove = msg.labels.toSet()
         msg.labels.clear()
-        val labelsToAdd = ctx.messageRepository.getLabels(Label.Type.TRASH)
+        val labelsToAdd = ctx.labelRepository.getLabels(Label.Type.TRASH)
         msg.addLabels(labelsToAdd)
         listener?.invoke(msg, labelsToAdd, labelsToRemove)
     }
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/ports/LabelRepository.kt b/core/src/main/kotlin/ch/dissem/bitmessage/ports/LabelRepository.kt
new file mode 100644
index 0000000..5484ca9
--- /dev/null
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/ports/LabelRepository.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright 2017 Christian Basler
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ch.dissem.bitmessage.ports
+
+import ch.dissem.bitmessage.entity.valueobject.Label
+
+interface LabelRepository {
+    fun getLabels(): List<Label>
+
+    fun getLabels(vararg types: Label.Type): List<Label>
+
+    fun save(label: Label)
+}
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/ports/MessageRepository.kt b/core/src/main/kotlin/ch/dissem/bitmessage/ports/MessageRepository.kt
index 68fac7a..272abae 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/ports/MessageRepository.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/ports/MessageRepository.kt
@@ -24,12 +24,6 @@ import ch.dissem.bitmessage.entity.valueobject.Label
 import java.util.*
 
 interface MessageRepository {
-    fun getLabels(): List<Label>
-
-    fun getLabels(vararg types: Label.Type): List<Label>
-
-    fun save(label: Label)
-
     fun countUnread(label: Label?): Int
 
     fun getAllMessages(): List<Plaintext>
diff --git a/core/src/main/kotlin/ch/dissem/bitmessage/utils/Collections.kt b/core/src/main/kotlin/ch/dissem/bitmessage/utils/Collections.kt
index 47626d2..57a7c1b 100644
--- a/core/src/main/kotlin/ch/dissem/bitmessage/utils/Collections.kt
+++ b/core/src/main/kotlin/ch/dissem/bitmessage/utils/Collections.kt
@@ -16,6 +16,7 @@
 
 package ch.dissem.bitmessage.utils
 
+import ch.dissem.bitmessage.exception.ApplicationException
 import java.util.*
 
 object Collections {
@@ -67,4 +68,12 @@ object Collections {
         }
         throw IllegalArgumentException("Empty collection? Size: " + collection.size)
     }
+
+    @JvmStatic fun <T> single(collection: Collection<T>): T? {
+        return when (collection.size) {
+            0 -> null
+            1 -> collection.iterator().next()
+            else -> throw ApplicationException("This shouldn't happen, found ${collection.size} items, one or none was expected")
+        }
+    }
 }
diff --git a/core/src/test/kotlin/ch/dissem/bitmessage/BitmessageContextTest.kt b/core/src/test/kotlin/ch/dissem/bitmessage/BitmessageContextTest.kt
index 28b9d43..dca4744 100644
--- a/core/src/test/kotlin/ch/dissem/bitmessage/BitmessageContextTest.kt
+++ b/core/src/test/kotlin/ch/dissem/bitmessage/BitmessageContextTest.kt
@@ -98,6 +98,7 @@ class BitmessageContextTest {
         cryptography = BouncyCryptography()
         inventory = testInventory
         listener = testListener
+        labelRepo = mock()
         messageRepo = mock()
         networkHandler = mock {
             on { getNetworkStatus() } doReturn Property("test", "mocked")
diff --git a/core/src/test/kotlin/ch/dissem/bitmessage/utils/TestUtils.kt b/core/src/test/kotlin/ch/dissem/bitmessage/utils/TestUtils.kt
index 0828869..fb80804 100644
--- a/core/src/test/kotlin/ch/dissem/bitmessage/utils/TestUtils.kt
+++ b/core/src/test/kotlin/ch/dissem/bitmessage/utils/TestUtils.kt
@@ -113,6 +113,7 @@ object TestUtils {
         nodeRegistry: NodeRegistry = mock {},
         networkHandler: NetworkHandler = mock {},
         addressRepository: AddressRepository = mock {},
+        labelRepository: LabelRepository = mock {},
         messageRepository: MessageRepository = mock {},
         proofOfWorkRepository: ProofOfWorkRepository = mock {},
         proofOfWorkEngine: ProofOfWorkEngine = mock {},
@@ -129,6 +130,7 @@ object TestUtils {
             nodeRegistry,
             networkHandler,
             addressRepository,
+            labelRepository,
             messageRepository,
             proofOfWorkRepository,
             proofOfWorkEngine,
diff --git a/demo/src/main/java/ch/dissem/bitmessage/demo/Application.java b/demo/src/main/java/ch/dissem/bitmessage/demo/Application.java
index d915215..50c2285 100644
--- a/demo/src/main/java/ch/dissem/bitmessage/demo/Application.java
+++ b/demo/src/main/java/ch/dissem/bitmessage/demo/Application.java
@@ -285,7 +285,7 @@ public class Application {
     }
 
     private void labels() {
-        List<Label> labels = ctx.messages().getLabels();
+        List<Label> labels = ctx.labels().getLabels();
         String command;
         do {
             System.out.println();
diff --git a/demo/src/test/java/ch/dissem/bitmessage/SystemTest.java b/demo/src/test/java/ch/dissem/bitmessage/SystemTest.java
index 6121f65..238a986 100644
--- a/demo/src/test/java/ch/dissem/bitmessage/SystemTest.java
+++ b/demo/src/test/java/ch/dissem/bitmessage/SystemTest.java
@@ -35,9 +35,7 @@ import java.util.concurrent.TimeUnit;
 
 import static ch.dissem.bitmessage.entity.payload.Pubkey.Feature.DOES_ACK;
 import static ch.dissem.bitmessage.utils.UnixTime.MINUTE;
-import static com.nhaarman.mockito_kotlin.MockitoKt.spy;
-import static com.nhaarman.mockito_kotlin.MockitoKt.timeout;
-import static com.nhaarman.mockito_kotlin.MockitoKt.verify;
+import static com.nhaarman.mockito_kotlin.MockitoKt.*;
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.junit.Assert.assertThat;
 import static org.mockito.ArgumentMatchers.any;
@@ -71,6 +69,7 @@ public class SystemTest {
             alice = new BitmessageContext.Builder()
                 .addressRepo(new JdbcAddressRepository(aliceDB))
                 .inventory(new JdbcInventory(aliceDB))
+                .labelRepo(new JdbcLabelRepository(aliceDB))
                 .messageRepo(new JdbcMessageRepository(aliceDB))
                 .powRepo(new JdbcProofOfWorkRepository(aliceDB))
                 .nodeRegistry(new TestNodeRegistry(bobPort))
@@ -89,6 +88,7 @@ public class SystemTest {
             bob = new BitmessageContext.Builder()
                 .addressRepo(new JdbcAddressRepository(bobDB))
                 .inventory(new JdbcInventory(bobDB))
+                .labelRepo(new JdbcLabelRepository(bobDB))
                 .messageRepo(new JdbcMessageRepository(bobDB))
                 .powRepo(new JdbcProofOfWorkRepository(bobDB))
                 .nodeRegistry(new TestNodeRegistry(alicePort))
diff --git a/networking/src/test/kotlin/ch/dissem/bitmessage/networking/NetworkHandlerTest.kt b/networking/src/test/kotlin/ch/dissem/bitmessage/networking/NetworkHandlerTest.kt
index 2931cfe..cd030f5 100644
--- a/networking/src/test/kotlin/ch/dissem/bitmessage/networking/NetworkHandlerTest.kt
+++ b/networking/src/test/kotlin/ch/dissem/bitmessage/networking/NetworkHandlerTest.kt
@@ -72,6 +72,7 @@ class NetworkHandlerTest {
             nodeRegistry = TestNodeRegistry()
             networkHandler = peerNetworkHandler
             addressRepo = mock()
+            labelRepo = mock()
             messageRepo = mock()
             proofOfWorkRepo = mock()
             customCommandHandler = object : CustomCommandHandler {
@@ -102,6 +103,7 @@ class NetworkHandlerTest {
             nodeRegistry = TestNodeRegistry(peerAddress)
             networkHandler = nodeNetworkHandler
             addressRepo = mock()
+            labelRepo = mock()
             messageRepo = mock()
             proofOfWorkRepo = mock()
             customCommandHandler = object : CustomCommandHandler {
diff --git a/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepository.kt b/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepository.kt
new file mode 100644
index 0000000..d7f2f3f
--- /dev/null
+++ b/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepository.kt
@@ -0,0 +1,123 @@
+/*
+ * Copyright 2017 Christian Basler
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ch.dissem.bitmessage.repository
+
+import ch.dissem.bitmessage.entity.valueobject.Label
+import ch.dissem.bitmessage.ports.AbstractLabelRepository
+import ch.dissem.bitmessage.ports.LabelRepository
+import org.slf4j.LoggerFactory
+import java.sql.Connection
+import java.sql.ResultSet
+import java.sql.SQLException
+import java.util.*
+
+class JdbcLabelRepository(private val config: JdbcConfig) : AbstractLabelRepository(), LabelRepository {
+
+    override fun find(where: String): List<Label> {
+        try {
+            config.getConnection().use { connection ->
+                return findLabels(connection, where)
+            }
+        } catch (e: SQLException) {
+            LOG.error(e.message, e)
+            return ArrayList()
+        }
+    }
+
+    override fun save(label: Label) {
+        config.getConnection().use { connection ->
+            if (label.id != null) {
+                connection.prepareStatement("UPDATE Label SET label=?, type=?, color=?, ord=? WHERE id=?").use { ps ->
+                    ps.setString(1, label.toString())
+                    ps.setString(2, label.type?.name)
+                    ps.setInt(3, label.color)
+                    ps.setInt(4, label.ord)
+                    ps.setInt(5, label.id as Int)
+                    ps.executeUpdate()
+                }
+            } else {
+                try {
+                    connection.autoCommit = false
+                    var exists = false
+                    connection.prepareStatement("SELECT COUNT(1) FROM Label WHERE label=?").use { ps ->
+                        ps.setString(1, label.toString())
+                        val rs = ps.executeQuery()
+                        if (rs.next()) {
+                            exists = rs.getInt(1) > 0
+                        }
+                    }
+
+                    if (exists) {
+                        connection.prepareStatement("UPDATE Label SET type=?, color=?, ord=? WHERE label=?").use { ps ->
+                            ps.setString(1, label.type?.name)
+                            ps.setInt(2, label.color)
+                            ps.setInt(3, label.ord)
+                            ps.setString(4, label.toString())
+                            ps.executeUpdate()
+                        }
+                    } else {
+                        connection.prepareStatement("INSERT INTO Label (label, type, color, ord) VALUES (?, ?, ?, ?)").use { ps ->
+                            ps.setString(1, label.toString())
+                            ps.setString(2, label.type?.name)
+                            ps.setInt(3, label.color)
+                            ps.setInt(4, label.ord)
+                            ps.executeUpdate()
+                        }
+                    }
+                    connection.commit()
+                } catch (e: Exception) {
+                    connection.rollback()
+                    throw e
+                }
+            }
+        }
+    }
+
+    private fun findLabels(connection: Connection, where: String): List<Label> {
+        val result = ArrayList<Label>()
+        try {
+            connection.createStatement().use { stmt ->
+                stmt.executeQuery("SELECT id, label, type, color, ord FROM Label WHERE $where").use { rs ->
+                    while (rs.next()) {
+                        result.add(getLabel(rs))
+                    }
+                }
+            }
+        } catch (e: SQLException) {
+            LOG.error(e.message, e)
+        }
+
+        return result
+    }
+
+    companion object {
+        private val LOG = LoggerFactory.getLogger(JdbcLabelRepository::class.java)
+
+        internal fun getLabel(rs: ResultSet): Label {
+            val typeName = rs.getString("type")
+            val type = if (typeName == null) {
+                null
+            } else {
+                Label.Type.valueOf(typeName)
+            }
+            val label = Label(rs.getString("label"), type, rs.getInt("color"), rs.getInt("ord"))
+            label.id = rs.getLong("id")
+
+            return label
+        }
+    }
+}
diff --git a/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepository.kt b/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepository.kt
index 6936d97..c7f6d3a 100644
--- a/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepository.kt
+++ b/repositories/src/main/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepository.kt
@@ -16,6 +16,7 @@
 
 package ch.dissem.bitmessage.repository
 
+import ch.dissem.bitmessage.entity.BitmessageAddress
 import ch.dissem.bitmessage.entity.Plaintext
 import ch.dissem.bitmessage.entity.valueobject.InventoryVector
 import ch.dissem.bitmessage.entity.valueobject.Label
@@ -29,79 +30,6 @@ import java.util.*
 
 class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRepository(), MessageRepository {
 
-    override fun findLabels(where: String): List<Label> {
-        try {
-            config.getConnection().use { connection ->
-                return findLabels(connection, where)
-            }
-        } catch (e: SQLException) {
-            LOG.error(e.message, e)
-            return ArrayList()
-        }
-    }
-
-    private fun getLabel(rs: ResultSet): Label {
-        val typeName = rs.getString("type")
-        val type = if (typeName == null) {
-            null
-        } else {
-            Label.Type.valueOf(typeName)
-        }
-        val label = Label(rs.getString("label"), type, rs.getInt("color"), rs.getInt("ord"))
-        label.id = rs.getLong("id")
-
-        return label
-    }
-
-    override fun save(label: Label) {
-        config.getConnection().use { connection ->
-            if (label.id != null) {
-                connection.prepareStatement("UPDATE Label SET label=?, type=?, color=?, ord=? WHERE id=?").use { ps ->
-                    ps.setString(1, label.toString())
-                    ps.setString(2, label.type?.name)
-                    ps.setInt(3, label.color)
-                    ps.setInt(4, label.ord)
-                    ps.setInt(5, label.id as Int)
-                    ps.executeUpdate()
-                }
-            } else {
-                try {
-                    connection.autoCommit = false
-                    var exists = false
-                    connection.prepareStatement("SELECT COUNT(1) FROM Label WHERE label=?").use { ps ->
-                        ps.setString(1, label.toString())
-                        val rs = ps.executeQuery()
-                        if (rs.next()) {
-                            exists = rs.getInt(1) > 0
-                        }
-                    }
-
-                    if (exists) {
-                        connection.prepareStatement("UPDATE Label SET type=?, color=?, ord=? WHERE label=?").use { ps ->
-                            ps.setString(1, label.type?.name)
-                            ps.setInt(2, label.color)
-                            ps.setInt(3, label.ord)
-                            ps.setString(4, label.toString())
-                            ps.executeUpdate()
-                        }
-                    } else {
-                        connection.prepareStatement("INSERT INTO Label (label, type, color, ord) VALUES (?, ?, ?, ?)").use { ps ->
-                            ps.setString(1, label.toString())
-                            ps.setString(2, label.type?.name)
-                            ps.setInt(3, label.color)
-                            ps.setInt(4, label.ord)
-                            ps.executeUpdate()
-                        }
-                    }
-                    connection.commit()
-                } catch (e: Exception) {
-                    connection.rollback()
-                    throw e
-                }
-            }
-        }
-    }
-
     override fun countUnread(label: Label?): Int {
         val where = if (label == null) {
             ""
@@ -136,26 +64,7 @@ class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRep
                         """SELECT id, iv, type, sender, recipient, data, ack_data, sent, received, initial_hash, status, ttl, retries, next_try, conversation
                            FROM Message WHERE $where $limit""").use { rs ->
                         while (rs.next()) {
-                            val iv = rs.getBytes("iv")
-                            val data = rs.getBinaryStream("data")
-                            val type = Plaintext.Type.valueOf(rs.getString("type"))
-                            val builder = Plaintext.readWithoutSignature(type, data)
-                            val id = rs.getLong("id")
-                            builder.id(id)
-                            builder.IV(InventoryVector.fromHash(iv))
-                            builder.from(ctx.addressRepository.getAddress(rs.getString("sender"))!!)
-                            rs.getString("recipient")?.let { builder.to(ctx.addressRepository.getAddress(it)) }
-                            builder.ackData(rs.getBytes("ack_data"))
-                            builder.sent(rs.getObject("sent") as Long?)
-                            builder.received(rs.getObject("received") as Long?)
-                            builder.status(Plaintext.Status.valueOf(rs.getString("status")))
-                            builder.ttl(rs.getLong("ttl"))
-                            builder.retries(rs.getInt("retries"))
-                            builder.nextTry(rs.getObject("next_try") as Long?)
-                            builder.conversation(rs.getObject("conversation") as UUID? ?: UUID.randomUUID())
-                            builder.labels(findLabels(connection,
-                                "id IN (SELECT label_id FROM Message_Label WHERE message_id=$id) ORDER BY ord"))
-                            val message = builder.build()
+                            val message = getMessage(connection, rs)
                             message.initialHash = rs.getBytes("initial_hash")
                             result.add(message)
                         }
@@ -165,17 +74,38 @@ class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRep
         } catch (e: SQLException) {
             LOG.error(e.message, e)
         }
-
         return result
     }
 
+    private fun getMessage(connection: Connection, rs: ResultSet): Plaintext {
+        return Plaintext.readWithoutSignature(
+            Plaintext.Type.valueOf(rs.getString("type")),
+            rs.getBinaryStream("data")
+        ).build {
+            id = rs.getLong("id")
+            inventoryVector = InventoryVector.fromHash(rs.getBytes("iv"))
+            from = rs.getString("sender")?.let { ctx.addressRepository.getAddress(it) ?: BitmessageAddress(it) }
+            to = rs.getString("recipient")?.let { ctx.addressRepository.getAddress(it) ?: BitmessageAddress(it) }
+            ackData = rs.getBytes("ack_data")
+            sent = rs.getObject("sent") as Long?
+            received = rs.getObject("received") as Long?
+            status = Plaintext.Status.valueOf(rs.getString("status"))
+            ttl = rs.getLong("ttl")
+            retries = rs.getInt("retries")
+            nextTry = rs.getObject("next_try") as Long?
+            conversation = rs.getObject("conversation") as UUID? ?: UUID.randomUUID()
+            labels = findLabels(connection,
+                "id IN (SELECT label_id FROM Message_Label WHERE message_id=$id) ORDER BY ord")
+        }
+    }
+
     private fun findLabels(connection: Connection, where: String): List<Label> {
         val result = ArrayList<Label>()
         try {
             connection.createStatement().use { stmt ->
                 stmt.executeQuery("SELECT id, label, type, color, ord FROM Label WHERE $where").use { rs ->
                     while (rs.next()) {
-                        result.add(getLabel(rs))
+                        result.add(JdbcLabelRepository.getLabel(rs))
                     }
                 }
             }
@@ -258,19 +188,7 @@ class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRep
                 "status, initial_hash, ttl, retries, next_try, conversation) " +
                 "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
             Statement.RETURN_GENERATED_KEYS).use { ps ->
-            ps.setBytes(1, if (message.inventoryVector == null) null else message.inventoryVector!!.hash)
-            ps.setString(2, message.type.name)
-            ps.setString(3, message.from.address)
-            ps.setString(4, if (message.to == null) null else message.to!!.address)
-            writeBlob(ps, 5, message)
-            ps.setBytes(6, message.ackData)
-            ps.setObject(7, message.sent)
-            ps.setObject(8, message.received)
-            ps.setString(9, message.status.name)
-            ps.setBytes(10, message.initialHash)
-            ps.setLong(11, message.ttl)
-            ps.setInt(12, message.retries)
-            ps.setObject(13, message.nextTry)
+            prepare(ps, message)
             ps.setObject(14, message.conversationId)
 
             try {
@@ -291,24 +209,29 @@ class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRep
             "UPDATE Message SET iv=?, type=?, sender=?, recipient=?, data=?, ack_data=?, sent=?, received=?, " +
                 "status=?, initial_hash=?, ttl=?, retries=?, next_try=? " +
                 "WHERE id=?").use { ps ->
-            ps.setBytes(1, if (message.inventoryVector == null) null else message.inventoryVector!!.hash)
-            ps.setString(2, message.type.name)
-            ps.setString(3, message.from.address)
-            ps.setString(4, if (message.to == null) null else message.to!!.address)
-            writeBlob(ps, 5, message)
-            ps.setBytes(6, message.ackData)
-            ps.setObject(7, message.sent)
-            ps.setObject(8, message.received)
-            ps.setString(9, message.status.name)
-            ps.setBytes(10, message.initialHash)
-            ps.setLong(11, message.ttl)
-            ps.setInt(12, message.retries)
-            ps.setObject(13, message.nextTry)
+            prepare(ps, message)
             ps.setLong(14, (message.id as Long?)!!)
             ps.executeUpdate()
         }
     }
 
+    private fun prepare(ps: PreparedStatement, message: Plaintext): Int{
+        ps.setBytes(1, if (message.inventoryVector == null) null else message.inventoryVector!!.hash)
+        ps.setString(2, message.type.name)
+        ps.setString(3, message.from.address)
+        ps.setString(4, if (message.to == null) null else message.to!!.address)
+        writeBlob(ps, 5, message)
+        ps.setBytes(6, message.ackData)
+        ps.setObject(7, message.sent)
+        ps.setObject(8, message.received)
+        ps.setString(9, message.status.name)
+        ps.setBytes(10, message.initialHash)
+        ps.setLong(11, message.ttl)
+        ps.setInt(12, message.retries)
+        ps.setObject(13, message.nextTry)
+        return 14
+    }
+
     override fun remove(message: Plaintext) {
         try {
             config.getConnection().use { connection ->
@@ -332,7 +255,6 @@ class JdbcMessageRepository(private val config: JdbcConfig) : AbstractMessageRep
         } catch (e: SQLException) {
             LOG.error(e.message, e)
         }
-
     }
 
     override fun findConversations(label: Label?): List<UUID> {
diff --git a/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepositoryTest.kt b/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepositoryTest.kt
new file mode 100644
index 0000000..f7822ec
--- /dev/null
+++ b/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcLabelRepositoryTest.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2017 Christian Basler
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package ch.dissem.bitmessage.repository
+
+import ch.dissem.bitmessage.entity.valueobject.Label
+import ch.dissem.bitmessage.ports.LabelRepository
+import org.junit.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+
+class JdbcLabelRepositoryTest : TestBase() {
+
+    private lateinit var repo: LabelRepository
+
+    @Before
+    fun setUp() {
+        val config = TestJdbcConfig()
+        config.reset()
+        repo = JdbcLabelRepository(config)
+    }
+
+    @Test
+    fun `ensure labels are retrieved`() {
+        val labels = repo.getLabels()
+        assertEquals(5, labels.size.toLong())
+    }
+
+    @Test
+    fun `ensure labels can be retrieved by type`() {
+        val labels = repo.getLabels(Label.Type.INBOX)
+        assertEquals(1, labels.size.toLong())
+        assertEquals("Inbox", labels[0].toString())
+    }
+
+}
diff --git a/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepositoryTest.kt b/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepositoryTest.kt
index 91753f5..8d078e0 100644
--- a/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepositoryTest.kt
+++ b/repositories/src/test/kotlin/ch/dissem/bitmessage/repository/JdbcMessageRepositoryTest.kt
@@ -26,6 +26,7 @@ import ch.dissem.bitmessage.entity.valueobject.ExtendedEncoding
 import ch.dissem.bitmessage.entity.valueobject.Label
 import ch.dissem.bitmessage.entity.valueobject.PrivateKey
 import ch.dissem.bitmessage.entity.valueobject.extended.Message
+import ch.dissem.bitmessage.ports.LabelRepository
 import ch.dissem.bitmessage.ports.MessageRepository
 import ch.dissem.bitmessage.utils.TestUtils
 import ch.dissem.bitmessage.utils.TestUtils.mockedInternalContext
@@ -46,6 +47,7 @@ class JdbcMessageRepositoryTest : TestBase() {
     private lateinit var identity: BitmessageAddress
 
     private lateinit var repo: MessageRepository
+    private lateinit var labelRepo: LabelRepository
 
     private lateinit var inbox: Label
     private lateinit var sent: Label
@@ -58,6 +60,7 @@ class JdbcMessageRepositoryTest : TestBase() {
         config.reset()
         val addressRepo = JdbcAddressRepository(config)
         repo = JdbcMessageRepository(config)
+        labelRepo = JdbcLabelRepository(config)
         mockedInternalContext(
             cryptography = BouncyCryptography(),
             addressRepository = addressRepo,
@@ -76,29 +79,16 @@ class JdbcMessageRepositoryTest : TestBase() {
         identity = BitmessageAddress(PrivateKey(false, 1, 1000, 1000, DOES_ACK))
         addressRepo.save(identity)
 
-        inbox = repo.getLabels(Label.Type.INBOX)[0]
-        sent = repo.getLabels(Label.Type.SENT)[0]
-        drafts = repo.getLabels(Label.Type.DRAFT)[0]
-        unread = repo.getLabels(Label.Type.UNREAD)[0]
+        inbox = labelRepo.getLabels(Label.Type.INBOX)[0]
+        sent = labelRepo.getLabels(Label.Type.SENT)[0]
+        drafts = labelRepo.getLabels(Label.Type.DRAFT)[0]
+        unread = labelRepo.getLabels(Label.Type.UNREAD)[0]
 
         addMessage(contactA, identity, Plaintext.Status.RECEIVED, inbox, unread)
         addMessage(identity, contactA, Plaintext.Status.DRAFT, drafts)
         addMessage(identity, contactB, Plaintext.Status.DRAFT, unread)
     }
 
-    @Test
-    fun `ensure labels are retrieved`() {
-        val labels = repo.getLabels()
-        assertEquals(5, labels.size.toLong())
-    }
-
-    @Test
-    fun `ensure labels can be retrieved by type`() {
-        val labels = repo.getLabels(Label.Type.INBOX)
-        assertEquals(1, labels.size.toLong())
-        assertEquals("Inbox", labels[0].toString())
-    }
-
     @Test
     fun `ensure messages can be found by label`() {
         val messages = repo.findMessages(inbox)
diff --git a/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifExporterTest.kt b/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifExporterTest.kt
index 523b4bb..2244ef8 100644
--- a/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifExporterTest.kt
+++ b/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifExporterTest.kt
@@ -53,6 +53,7 @@ class WifExporterTest {
             cryptography = BouncyCryptography()
             networkHandler = mock()
             inventory = mock()
+            labelRepo = mock()
             messageRepo = mock()
             proofOfWorkRepo = mock()
             nodeRegistry = mock()
diff --git a/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifImporterTest.kt b/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifImporterTest.kt
index 0be2bf7..758bb39 100644
--- a/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifImporterTest.kt
+++ b/wif/src/test/kotlin/ch/dissem/bitmessage/wif/WifImporterTest.kt
@@ -54,6 +54,7 @@ class WifImporterTest {
             cryptography = BouncyCryptography()
             networkHandler = mock()
             inventory = mock()
+            labelRepo = mock()
             messageRepo = mock()
             proofOfWorkRepo = mock()
             nodeRegistry = mock()