org.jetbrains.kotlinx.jupyter.test.protocol.KernelServerTestsBase.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of kotlin-jupyter-script-classpath-shadowed Show documentation

Kotlin Jupyter kernel script classpath with all dependencies inside one artifact

There is a newer version: 0.12.0-356

Show newest version

package org.jetbrains.kotlinx.jupyter.test.protocol

import io.kotest.matchers.shouldBe
import io.kotest.matchers.types.shouldBeTypeOf
import org.jetbrains.kotlinx.jupyter.messaging.AbstractMessageContent
import org.jetbrains.kotlinx.jupyter.messaging.KernelStatus
import org.jetbrains.kotlinx.jupyter.messaging.Message
import org.jetbrains.kotlinx.jupyter.messaging.MessageData
import org.jetbrains.kotlinx.jupyter.messaging.MessageType
import org.jetbrains.kotlinx.jupyter.messaging.StatusReply
import org.jetbrains.kotlinx.jupyter.messaging.makeHeader
import org.jetbrains.kotlinx.jupyter.messaging.sendMessage
import org.jetbrains.kotlinx.jupyter.messaging.toMessage
import org.jetbrains.kotlinx.jupyter.protocol.JupyterSocketBase
import org.jetbrains.kotlinx.jupyter.protocol.JupyterSocketInfo
import org.jetbrains.kotlinx.jupyter.protocol.JupyterSocketSide
import org.jetbrains.kotlinx.jupyter.protocol.SocketWrapper
import org.jetbrains.kotlinx.jupyter.protocol.createSocket
import org.jetbrains.kotlinx.jupyter.startup.createKotlinKernelConfig
import org.jetbrains.kotlinx.jupyter.startup.createRandomKernelPorts
import org.jetbrains.kotlinx.jupyter.test.classpath
import org.jetbrains.kotlinx.jupyter.test.testLoggerFactory
import org.junit.jupiter.api.AfterEach
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.TestInfo
import org.zeromq.ZMQ
import java.io.File
import java.util.UUID

/**
 * Base class for tests that need to have fine-tuned control over the kernel server execution.
 *
 * @param runServerInSeparateProcess The debugger only works if the kernel is running in the same process
 * as the test. The default is to run in a separate one (similar to how it would work in
 * production). To enable debugging, set this value to `false`
 */
abstract class KernelServerTestsBase(protected val runServerInSeparateProcess: Boolean) {
    protected abstract val context: ZMQ.Context

    protected val kernelConfig =
        createKotlinKernelConfig(
            ports = createRandomKernelPorts(),
            signatureKey = "abc",
            scriptClasspath = classpath,
            homeDir = File(""),
        )

    private val sessionId = UUID.randomUUID().toString()
    private val messageId = listOf(byteArrayOf(1))

    private val executor = if (runServerInSeparateProcess) ProcessServerTestExecutor() else ThreadServerTestExecutor()

    open fun beforeEach() {}

    open fun afterEach() {}

    @BeforeEach
    fun setupServer(testInfo: TestInfo) {
        executor.setUp(testInfo, kernelConfig)
        beforeEach()
    }

    @AfterEach
    fun teardownServer() {
        afterEach()
        executor.tearDown()
    }

    fun createClientSocket(socketInfo: JupyterSocketInfo) =
        createSocket(
            testLoggerFactory,
            socketInfo,
            context,
            kernelConfig,
            JupyterSocketSide.CLIENT,
        )

    fun JupyterSocketBase.sendMessage(
        msgType: MessageType,
        content: AbstractMessageContent?,
    ) {
        sendMessage(Message(id = messageId, MessageData(header = makeHeader(msgType, sessionId = sessionId), content = content)))
    }

    fun JupyterSocketBase.receiveMessage() = receiveRawMessage()!!.toMessage()

    fun JupyterSocketBase.receiveStatusReply(): StatusReply {
        (this as? SocketWrapper)?.name shouldBe JupyterSocketInfo.IOPUB.name
        receiveMessage().apply {
            return content.shouldBeTypeOf()
        }
    }

    inline fun JupyterSocketBase.wrapActionInBusyIdleStatusChange(action: () -> Unit) {
        receiveStatusReply().status shouldBe KernelStatus.BUSY
        action()
        receiveStatusReply().status shouldBe KernelStatus.IDLE
    }
}