Add helper methods to support running misk service tests in

parallel

GitOrigin-RevId: cd048b38737715800dfdcb21650d3a8daa636f25

Author: mmollaverdi

misk-jdbc/api/misk-jdbc.api

@@ -590,6 +590,11 @@ public final class misk/jdbc/MySqlTestDatabasePoolBackend : misk/jdbc/TestDataba
 	public fun showDatabases ()Ljava/util/Set;
 }
 
+public final class misk/jdbc/ParitionedDataSourceConfigKt {
+	public static final fun updateForParallelTests (Lmisk/jdbc/DataSourceClusterConfig;)Lmisk/jdbc/DataSourceClusterConfig;
+	public static final fun updateForParallelTests (Lmisk/jdbc/DataSourceConfig;)Lmisk/jdbc/DataSourceConfig;
+}
+
 public final class misk/jdbc/PingDatabaseService : com/google/common/util/concurrent/AbstractIdleService {
 	public fun <init> (Lmisk/jdbc/DataSourceConfig;Lwisp/deployment/Deployment;)V
 }

misk-jdbc/build.gradle.kts

@@ -56,6 +56,7 @@ dependencies {
   testFixturesImplementation(project(":misk"))
   testFixturesImplementation(project(":misk-core"))
   testFixturesImplementation(project(":misk-service"))
+  testFixturesImplementation(project(":misk-testing"))
   testFixturesRuntimeOnly(libs.hsqldb)
 
   testImplementation(libs.assertj)

misk-jdbc/src/testFixtures/kotlin/misk/jdbc/ParitionedDataSourceConfig.kt

@@ -0,0 +1,12 @@
+package misk.jdbc
+
+import misk.testing.updateForParallelTests
+
+fun DataSourceConfig.updateForParallelTests(): DataSourceConfig = this.updateForParallelTests { config, partitionId ->
+  config.copy(database = "${config.database}_$partitionId")
+}
+
+fun DataSourceClusterConfig.updateForParallelTests(): DataSourceClusterConfig = this.copy(
+  writer = writer.updateForParallelTests(),
+  reader = reader?.updateForParallelTests()
+)

misk-testing/api/misk-testing.api

@@ -222,6 +222,10 @@ public final class misk/testing/MockTracingBackendModule : misk/inject/KAbstract
 	public fun <init> ()V
 }
 
+public final class misk/testing/ParallelTestsKt {
+	public static final fun updateForParallelTests (Ljava/lang/Object;Lkotlin/jvm/functions/Function2;)Ljava/lang/Object;
+}
+
 public final class misk/testing/TemporaryFolder {
 	public fun <init> (Ljava/nio/file/Path;)V
 	public final fun delete ()V

misk-testing/src/main/kotlin/misk/testing/ParallelTests.kt

@@ -0,0 +1,50 @@
+package misk.testing
+
+/**
+ * Provides helper methods for supporting tests running in parallel.
+ */
+internal object ParallelTests {
+
+  /**
+   * Can be used for determining if tests are running in parallel (with Gradle's `maxParallelForks`),
+   * and returns a unique numeric partition ID for the current test process.
+   *
+   * This assumes that the environment variable `MAX_TEST_PARALLEL_FORKS` is set to the number of `maxParallelForks`.
+   */
+  internal fun partitioned(): PartitionedTest {
+    val maxTestPartitions = System.getenv("MAX_TEST_PARALLEL_FORKS")?.toInt()
+    return if (maxTestPartitions != null && maxTestPartitions > 1) {
+      val testWorkerId = System.getProperty("org.gradle.test.worker", "0").toInt()
+      val partitionId = ((testWorkerId % maxTestPartitions) + 1)
+      PartitionedTest.Partitioned(partitionId)
+    } else {
+      PartitionedTest.NotPartitioned
+    }
+  }
+}
+
+internal sealed interface PartitionedTest {
+  data object NotPartitioned : PartitionedTest
+  data class Partitioned(val partitionId: Int) : PartitionedTest
+}
+
+/**
+ * Applies the `update` function to the given `T`, if the tests are running in parallel
+ * (with Gradle's `maxParallelForks`). The `update` function is also passed the partition ID.
+ *
+ * This can be used for updating configurations for the purpose of providing isolation between
+ * tests running across different parallel processes, such as appending the partition ID to the database name,
+ * so that it's unique for each process.
+ *
+ * This assumes that the environment variable `MAX_TEST_PARALLEL_FORKS` is set to the number of `maxParallelForks`.
+ */
+fun <T> T.updateForParallelTests(update: (T, Int) -> T): T {
+  return when (val partitionedTest = ParallelTests.partitioned()) {
+    is PartitionedTest.Partitioned -> {
+      update(this, partitionedTest.partitionId)
+    }
+
+    is PartitionedTest.NotPartitioned -> this
+  }
+}
+