Initial commit

.github/workflows/release.yml

@@ -0,0 +1,20 @@
+name: Release
+on:
+  push:
+    branches: [master, main]
+    tags: ["*"]
+jobs:
+  publish:
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/[email protected]
+        with:
+          fetch-depth: 0
+      - uses: olafurpg/setup-scala@v13
+      - uses: olafurpg/setup-gpg@v3
+      - run: sbt ci-release
+        env:
+          PGP_PASSPHRASE: ${{ secrets.PGP_PASSPHRASE }}
+          PGP_SECRET: ${{ secrets.PGP_SECRET }}
+          SONATYPE_PASSWORD: ${{ secrets.SONATYPE_PASSWORD }}
+          SONATYPE_USERNAME: ${{ secrets.SONATYPE_USERNAME }}

.github/workflows/scala.yml

@@ -0,0 +1,26 @@
+name: Scala CI
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+jobs:
+  jvm:
+    strategy:
+      fail-fast: false
+      matrix:
+        scala: [2.12.16]
+        java: [[email protected], [email protected]]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Set up environment
+        uses: olafurpg/setup-scala@v10
+        with:
+          java-version: ${{ matrix.java }}
+
+      - name: Run tests
+        run: sbt ++${{ matrix.scala}} fmtCheck test

.gitignore

@@ -0,0 +1,28 @@
+target
+.idea
+.idea_modules
+.bloop
+.bsp
+.metals
+/.classpath
+/.project
+/.settings
+/RUNNING_PID
+/out/
+*.iws
+*.iml
+/db
+.eclipse
+/lib/
+/logs/
+/modules
+tmp/
+test-result
+server.pid
+*.eml
+/dist/
+.cache
+/reference
+local.conf
+/logs
+publish.sbt

.scalafmt.conf

@@ -0,0 +1,18 @@
+version = "3.2.1"
+
+runner.dialect = scala213
+
+maxColumn = 120
+align.preset = most
+continuationIndent.defnSite = 2
+assumeStandardLibraryStripMargin = true
+docstrings.style = SpaceAsterisk
+lineEndings = preserve
+includeCurlyBraceInSelectChains = false
+danglingParentheses.preset = true
+spaces {
+  inImportCurlyBraces = true
+}
+optIn.annotationNewlines = true
+
+rewrite.rules = [SortImports, RedundantBraces]

README.md

@@ -1 +1,3 @@
-# sbt-test-shards
+# SBT Test Shards
+
+*An SBT plugin for splitting tests across multiple shards to speed up tests.*

build.sbt

@@ -0,0 +1,51 @@
+import sbtwelcome._
+
+inThisBuild(
+  List(
+    organization := "com.github.reibitto",
+    homepage     := Some(url("https://github.com/reibitto/sbt-test-shards")),
+    licenses     := List("Apache-2.0" -> url("https://www.apache.org/licenses/LICENSE-2.0")),
+    developers   := List(
+      Developer("reibitto", "reibitto", "[email protected]", url("https://reibitto.github.io"))
+    )
+  )
+)
+
+lazy val root = (project in file(".")).settings(
+  name         := "sbt-test-shards",
+  organization := "com.github.reibitto",
+  scalaVersion := "2.12.16",
+  sbtPlugin    := true
+)
+
+addCommandAlias("fmt", "all root/scalafmtSbt root/scalafmtAll")
+addCommandAlias("fmtCheck", "all root/scalafmtSbtCheck root/scalafmtCheckAll")
+
+logo :=
+  s"""
+     |                              ______ _____
+     |                       __________  /___  /_
+     |                       __  ___/_  __ \\  __/
+     |                       _(__  )_  /_/ / /_
+     |                       /____/ /_.___/\\__/
+     |  _____            _____           ______               _________
+     |  __  /______________  /_   __________  /_______ _____________  /_______
+     |  _  __/  _ \\_  ___/  __/   __  ___/_  __ \\  __ `/_  ___/  __  /__  ___/
+     |  / /_ /  __/(__  )/ /_     _(__  )_  / / / /_/ /_  /   / /_/ / _(__  )
+     |  \\__/ \\___//____/ \\__/     /____/ /_/ /_/\\__,_/ /_/    \\__,_/  /____/
+     |
+     |${version.value}
+     |
+     |${scala.Console.YELLOW}Scala ${scalaVersion.value}${scala.Console.RESET}
+     |
+     |""".stripMargin
+
+usefulTasks := Seq(
+  UsefulTask("a", "~compile", "Compile with file-watch enabled"),
+  UsefulTask("b", "fmt", "Run scalafmt on the entire project"),
+  UsefulTask("c", "publishLocal", "Publish the sbt plugin locally so that you can consume it from a different project")
+)
+
+logoColor := scala.Console.MAGENTA
+
+ThisBuild / organization := "com.github.reibitto"

project/build.properties

@@ -0,0 +1 @@
+sbt.version=1.7.1

project/plugins.sbt

@@ -0,0 +1,3 @@
+addSbtPlugin("org.scalameta"       % "sbt-scalafmt"   % "2.4.6")
+addSbtPlugin("com.github.sbt"      % "sbt-ci-release" % "1.5.10")
+addSbtPlugin("com.github.reibitto" % "sbt-welcome"    % "0.2.2")

src/main/scala/sbttestshards/ShardContext.scala

@@ -0,0 +1,5 @@
+package sbttestshards
+
+import sbt.Logger
+
+final case class ShardContext(testShard: Int, testShardCount: Int, logger: Logger)

src/main/scala/sbttestshards/ShardingAlgorithm.scala

@@ -0,0 +1,79 @@
+package sbttestshards
+
+import java.time.Duration
+
+// This trait is open so that users can implement a custom `ShardingAlgorithm` if they'd like
+trait ShardingAlgorithm {
+  def isInShard(specName: String, shardContext: ShardContext): Boolean
+}
+
+object ShardingAlgorithm {
+  final case object Always extends ShardingAlgorithm {
+    override def isInShard(specName: String, shardContext: ShardContext): Boolean = true
+  }
+
+  final case object Never extends ShardingAlgorithm {
+    override def isInShard(specName: String, shardContext: ShardContext): Boolean = false
+  }
+
+  final case object SuiteName extends ShardingAlgorithm {
+    override def isInShard(specName: String, shardContext: ShardContext): Boolean = {
+      val shouldRun = specName.hashCode % shardContext.testShardCount == shardContext.testShard
+
+      println(s"${specName} will run? ${shouldRun}")
+
+      shouldRun
+    }
+  }
+
+  final case class Balance(
+    tests: List[TestSuiteInfo],
+    bucketCount: Int,
+    fallbackShardingAlgorithm: ShardingAlgorithm = ShardingAlgorithm.SuiteName
+  ) extends ShardingAlgorithm {
+    // TODO: Median might be better here?
+    private val averageTime: Option[Duration] = {
+      val allTimeTaken = tests.flatMap(_.timeTaken)
+      allTimeTaken.reduceOption(_.plus(_)).map { d =>
+        if (d.isZero) Duration.ZERO
+        else d.dividedBy(allTimeTaken.length)
+      }
+    }
+
+    private final case class TestSuiteInfoSimple(name: String, timeTaken: Duration)
+    private final case class TestBucket(var tests: List[TestSuiteInfoSimple], var sum: Duration)
+
+    private def createBucketMap(testShardCount: Int) = {
+      val durationOrdering: Ordering[Duration] = (a: Duration, b: Duration) => a.compareTo(b)
+
+      val allTests = tests
+        .map(t => TestSuiteInfoSimple(t.name, t.timeTaken.getOrElse(averageTime.getOrElse(Duration.ZERO))))
+        .sortBy(_.timeTaken)(durationOrdering.reverse)
+
+      val buckets = Array.fill(testShardCount)(TestBucket(Nil, Duration.ZERO))
+
+      allTests.foreach { test =>
+        val minBucket = buckets.minBy(_.sum)
+
+        minBucket.tests = test :: minBucket.tests
+        minBucket.sum = minBucket.sum.plus(test.timeTaken)
+      }
+
+      buckets.zipWithIndex.flatMap { case (bucket, i) =>
+        bucket.tests.map { info =>
+          info.name -> i
+        }
+      }.toMap
+    }
+
+    // `bucketCount` doesn't necessary need to match `testShardCount`, but ideally it should be a multiple of it.
+    // TODO: Maybe print a warning if it's not a multiple of it.
+    private val bucketMap: Map[String, Int] = createBucketMap(bucketCount)
+
+    def isInShard(specName: String, shardContext: ShardContext): Boolean =
+      bucketMap.get(specName) match {
+        case Some(bucketIndex) => bucketIndex == shardContext.testShard
+        case None              => fallbackShardingAlgorithm.isInShard(specName, shardContext)
+      }
+  }
+}

src/main/scala/sbttestshards/TestBucketItem.scala

@@ -0,0 +1,5 @@
+package sbttestshards
+
+import java.time.Duration
+
+final case class TestBucketItem(name: String, timeTaken: Duration)

src/main/scala/sbttestshards/TestShardsPlugin.scala

@@ -0,0 +1,28 @@
+package sbttestshards
+
+import sbt.Keys.*
+import sbt.*
+
+object TestShardsPlugin extends AutoPlugin {
+  object autoImport {
+    val testShard         = settingKey[Int]("testShard")
+    val testShardCount    = settingKey[Int]("testShardCount")
+    val shardingAlgorithm = settingKey[ShardingAlgorithm]("shardingAlgorithm")
+  }
+
+  import autoImport.*
+
+  override def trigger = allRequirements
+
+  override lazy val projectSettings: Seq[Def.Setting[?]] =
+    Seq(
+      testShard         := 0,
+      testShardCount    := 1,
+      shardingAlgorithm := ShardingAlgorithm.SuiteName,
+      Test / testOptions += {
+        val shardContext = ShardContext(testShardCount.value, testShard.value, sLog.value)
+        Tests.Filter(specName => shardingAlgorithm.value.isInShard(specName, shardContext))
+      }
+    )
+
+}

src/main/scala/sbttestshards/TestSuiteInfo.scala

@@ -0,0 +1,5 @@
+package sbttestshards
+
+import java.time.Duration
+
+final case class TestSuiteInfo(name: String, timeTaken: Option[Duration])