ExhaustiveTestRunner.kt

annotation class Exhaustive(
  val sealedType: KClass<*>,
)

@Repeatable
annotation class ExhaustiveCase(
  val sealedType: KClass<*>,
  val test: KClass<*>,
)

class ExhaustiveTestRunner(testClass: Class<*>) : Suite(testClass, getTests(testClass)) {

  companion object {

    private fun getTests(testClass: Class<*>): Array<Class<*>> {
      check(testClass.isAnnotationPresent(Exhaustive::class.java)) {
        "ExhaustiveTestRunner's test class must be annotated with @Exhaustive"
      }

      val exhaustive = testClass.getAnnotation(Exhaustive::class.java)!!
      val sealedType = exhaustive.sealedType
      check(sealedType.isSealed) {
        "@Exhaustive class parameter must represent a sealed type"
      }

      val cases = testClass.getAnnotationsByType(ExhaustiveCase::class.java)
      check(cases.contentEquals(cases.sortedBy { it.sealedType.simpleName }.toTypedArray())) {
        "ExhaustiveCase must be in alphabetical order"
      }

      check(cases.contentEquals(cases.distinctBy { it.sealedType }.toTypedArray())) {
        "ExhaustiveCase must be distinct by key"
      }

      check(cases.contentEquals(cases.distinctBy { it.test }.toTypedArray())) {
        "ExhaustiveCase must be distinct by value"
      }

      val testLookup = cases.associate { it.sealedType to it.test }
      testLookup.keys.forEach { key ->
        check(key.isSubclassOf(sealedType)) {
          "Expected ExhaustiveCase key to be a subtype of $sealedType"
        }
      }

      return sealedType.sealedSubclasses
        .map { sealedSubclass ->
          checkNotNull(testLookup[sealedSubclass]) {
            "Missing required exhaustive case for type $sealedSubclass"
          }
        }
        .map(KClass<*>::java)
        .toTypedArray()
    }
  }
}