Store info about ZMQ socket types in one place

Author: ileasile

src/main/kotlin/org/jetbrains/kotlin/jupyter/config.kt

@@ -9,6 +9,7 @@ import kotlinx.coroutines.async
 import org.apache.commons.io.FileUtils
 import org.json.JSONObject
 import org.slf4j.LoggerFactory
+import org.zeromq.ZMQ
 import java.io.File
 import java.nio.file.Files
 import java.nio.file.Paths
@@ -30,12 +31,12 @@ val LibraryPropertiesFile = ".properties"
 
 internal val log by lazy { LoggerFactory.getLogger("ikotlin") }
 
-enum class JupyterSockets {
-    hb,
-    shell,
-    control,
-    stdin,
-    iopub
+enum class JupyterSockets(val zmqKernelType: Int, val zmqClientType: Int) {
+    hb(ZMQ.REP, ZMQ.REP),
+    shell(ZMQ.ROUTER, ZMQ.REQ),
+    control(ZMQ.ROUTER, ZMQ.REQ),
+    stdin(ZMQ.ROUTER, ZMQ.DEALER),
+    iopub(ZMQ.PUB, ZMQ.SUB)
 }
 
 data class OutputConfig(

src/main/kotlin/org/jetbrains/kotlin/jupyter/connection.kt

@@ -13,7 +13,7 @@ import javax.crypto.spec.SecretKeySpec
 
 class JupyterConnection(val config: KernelConfig): Closeable {
 
-    inner class Socket(private val socket: JupyterSockets, type: Int) : ZMQ.Socket(context, type) {
+    inner class Socket(private val socket: JupyterSockets, type: Int = socket.zmqKernelType) : ZMQ.Socket(context, type) {
         val name: String get() = socket.name
         init {
             val port = config.ports[socket.ordinal]
@@ -121,11 +121,11 @@ class JupyterConnection(val config: KernelConfig): Closeable {
 
     private val disposable = Disposable { }
 
-    val heartbeat = Socket(JupyterSockets.hb, ZMQ.REP)
-    val shell = Socket(JupyterSockets.shell, ZMQ.ROUTER)
-    val control = Socket(JupyterSockets.control, ZMQ.ROUTER)
-    val stdin = Socket(JupyterSockets.stdin, ZMQ.ROUTER)
-    val iopub = Socket(JupyterSockets.iopub, ZMQ.PUB)
+    val heartbeat = Socket(JupyterSockets.hb)
+    val shell = Socket(JupyterSockets.shell)
+    val control = Socket(JupyterSockets.control)
+    val stdin = Socket(JupyterSockets.stdin)
+    val iopub = Socket(JupyterSockets.iopub)
 
     val stdinIn = StdinInputStream()
 

src/test/kotlin/org/jetbrains/kotlin/jupyter/test/kernelServerTestsBase.kt

@@ -37,6 +37,10 @@ open class KernelServerTestsBase {
         server?.interrupt()
     }
 
+    inner class ClientSocket(context: ZMQ.Context, private val socket: JupyterSockets) : ZMQ.Socket(context, socket.zmqClientType) {
+        fun connect() = connect("${config.transport}://*:${config.ports[socket.ordinal]}")
+    }
+
     fun ZMQ.Socket.sendMessage(msgType: String, content : JsonObject) {
         sendMessage(Message(id = messageId, header = makeHeader(msgType), content = content), hmac)
     }