From ae854ee75b98cbb7334bb1b0c52e4c39525f2891 Mon Sep 17 00:00:00 2001
From: Bernhard Geisberger <bernhard.geisberger@jetbrains.com>
Date: Thu, 3 Jul 2025 23:32:23 +0200
Subject: [PATCH] [python, tests] tests for tcp socket transport from test
 scripts

GitOrigin-RevId: 499ed7e0d87e294c52ae8d95d0f456d3eb772c21
---
 .../env/pytest/test_keyboard_interrupt.py     |   4 +
 .../testing/PyTestScriptViaTcpSocketTest.kt   | 127 ++++++++++++++++++
 2 files changed, 131 insertions(+)
 create mode 100644 python/testData/testRunner/env/pytest/test_keyboard_interrupt.py
 create mode 100644 python/testSrc/com/jetbrains/env/python/testing/PyTestScriptViaTcpSocketTest.kt

diff --git a/python/testData/testRunner/env/pytest/test_keyboard_interrupt.py b/python/testData/testRunner/env/pytest/test_keyboard_interrupt.py
new file mode 100644
index 000000000000..8ab24b6a9b64
--- /dev/null
+++ b/python/testData/testRunner/env/pytest/test_keyboard_interrupt.py
@@ -0,0 +1,4 @@
+import pytest
+
+def test_keyboard_interrupt():
+    raise KeyboardInterrupt()
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/env/python/testing/PyTestScriptViaTcpSocketTest.kt b/python/testSrc/com/jetbrains/env/python/testing/PyTestScriptViaTcpSocketTest.kt
new file mode 100644
index 000000000000..8b02293485ca
--- /dev/null
+++ b/python/testSrc/com/jetbrains/env/python/testing/PyTestScriptViaTcpSocketTest.kt
@@ -0,0 +1,127 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.env.python.testing
+
+import com.intellij.execution.testframework.sm.ServiceMessageUtil
+import com.jetbrains.env.EnvTestTagsRequired
+import com.jetbrains.env.PyEnvTestCase
+import com.jetbrains.env.PyProcessWithConsoleTestTask
+import com.jetbrains.env.ut.PyTestTestProcessRunner
+import com.jetbrains.python.testing.PyTestConfiguration
+import com.jetbrains.python.tools.sdkTools.SdkCreationType
+import jetbrains.buildServer.messages.serviceMessages.ServiceMessage
+import org.junit.Assert
+import org.junit.Test
+import java.net.ServerSocket
+
+@EnvTestTagsRequired(tags = ["pytest"])
+class PyTestScriptViaTcpSocketTest : PyEnvTestCase() {
+
+  @Test
+  fun testReceivedMessageTypes() {
+    val messageReceiver = TcpSocketMessageReceiver()
+
+    runPythonTest(object : PyProcessWithConsoleTestTask<PyTestTestProcessRunner?>("/testRunner/env/pytest", SdkCreationType.EMPTY_SDK) {
+      protected override fun createProcessRunner(): PyTestTestProcessRunner {
+        return object : PyTestTestProcessRunner("test2.py", 0) {
+          override fun configurationCreatedAndWillLaunch(configuration: PyTestConfiguration) {
+            configuration.envs = configuration.envs + mapOf(
+              "JB_VERBOSE" to "true",
+              "JB_DISABLE_BUFFERING" to "true",
+              "JB_TEAMCITY_SOCKET_PORT" to messageReceiver.port.toString()
+            )
+            super.configurationCreatedAndWillLaunch(configuration)
+          }
+        }
+      }
+
+      override fun checkTestResults(
+        runner: PyTestTestProcessRunner,
+        stdout: String,
+        stderr: String,
+        all: String,
+        exitCode: Int,
+      ) {
+        Assert.assertEquals(runner.formattedTestTree, 0, runner.allTestsCount)
+
+        messageReceiver.join()
+
+        // with buffering disabled, we want to receive exactly these message types via socket
+        Assert.assertEquals(
+          messageReceiver.messages.map { it.messageName }.toSet(),
+          setOf("enteredTheMatrix", "testCount", "testSuiteStarted", "testStarted", "testFailed", "testFinished", "testSuiteFinished")
+        )
+      }
+    })
+  }
+
+  @Test
+  fun testKeyboardInterrupt() {
+    val messageReceiver = TcpSocketMessageReceiver()
+
+    runPythonTest(object : PyProcessWithConsoleTestTask<PyTestTestProcessRunner?>("/testRunner/env/pytest", SdkCreationType.EMPTY_SDK) {
+      protected override fun createProcessRunner(): PyTestTestProcessRunner {
+        return object : PyTestTestProcessRunner("test_keyboard_interrupt.py", 0) {
+          override fun configurationCreatedAndWillLaunch(configuration: PyTestConfiguration) {
+            configuration.envs = configuration.envs + mapOf(
+              "JB_VERBOSE" to "true",
+              "JB_DISABLE_BUFFERING" to "true",
+              "JB_TEAMCITY_SOCKET_PORT" to messageReceiver.port.toString()
+            )
+            super.configurationCreatedAndWillLaunch(configuration)
+          }
+        }
+      }
+
+      override fun checkTestResults(
+        runner: PyTestTestProcessRunner,
+        stdout: String,
+        stderr: String,
+        all: String,
+        exitCode: Int,
+      ) {
+        Assert.assertEquals(runner.formattedTestTree, 0, runner.allTestsCount)
+
+        messageReceiver.join()
+
+        val ignoredMessage = messageReceiver.messages.firstOrNull { it.messageName == "testIgnored" }
+        Assert.assertNotNull("No ignored message received", ignoredMessage)
+        Assert.assertEquals("Missing stopped information", "true", ignoredMessage!!.attributes["stopped"])
+      }
+    })
+  }
+}
+
+private class TcpSocketMessageReceiver {
+  private val serverSocket = ServerSocket(0)
+  private val receivedMessages = mutableListOf<ServiceMessage>()
+  private val socketThread: Thread = Thread {
+    serverSocket.use {
+      serverSocket.accept().use { socket ->
+        socket.inputStream.bufferedReader().use { reader ->
+          while (true) {
+            val line = reader.readLine() ?: break
+            val message = ServiceMessageUtil.parse(line, true)
+            Assert.assertNotNull("Failed to parse message: $line", message)
+            if (message != null) {
+              receivedMessages.add(message)
+            }
+          }
+        }
+      }
+    }
+  }
+
+  val port: Int
+    get() = serverSocket.localPort
+
+  val messages: List<ServiceMessage>
+    get() = receivedMessages.toList()
+
+  init {
+    socketThread.start()
+  }
+
+  fun join() {
+    socketThread.join()
+  }
+}
\ No newline at end of file