withEdgecases should keep shrinker #3812

kotest-property/src/commonMain/kotlin/io/kotest/property/arbitrary/edgecases.kt

@@ -1,6 +1,7 @@
 package io.kotest.property.arbitrary
 
 import io.kotest.property.Arb
+import io.kotest.property.Classifier
 import io.kotest.property.EdgeConfig
 import io.kotest.property.RandomSource
 import io.kotest.property.Sample
@@ -33,7 +34,11 @@ fun <A> Arb<A>.edgecases(iterations: Int = 100, rs: RandomSource = RandomSource.
 /**
  * Returns a new [Arb] with the supplied edge cases replacing any existing edge cases.
  */
-fun <A> Arb<A>.withEdgecases(edgecases: List<A>): Arb<A> = arbitrary(edgecases) { this@withEdgecases.next(it) }
+fun <A> Arb<A>.withEdgecases(edgecases: List<A>): Arb<A> = object : Arb<A>() {
+   override fun edgecase(rs: RandomSource): A? = if (edgecases.isEmpty()) null else edgecases.random(rs.random)
+   override fun sample(rs: RandomSource): Sample<A> = this@withEdgecases.sample(rs)
+   override val classifier: Classifier<out A>? = this@withEdgecases.classifier
+}
 
 /**
  * Returns a new [Arb] with the supplied edge cases replacing any existing edge cases.

kotest-property/src/jvmTest/kotlin/com/sksamuel/kotest/property/shrinking/ShrinkingTest.kt

@@ -5,7 +5,10 @@ import io.kotest.core.spec.style.FunSpec
 import io.kotest.extensions.system.captureStandardOut
 import io.kotest.matchers.ints.shouldBeLessThan
 import io.kotest.matchers.string.shouldContain
+import io.kotest.property.Arb
 import io.kotest.property.PropTestConfig
+import io.kotest.property.arbitrary.int
+import io.kotest.property.arbitrary.withEdgecases
 import io.kotest.property.checkAll
 
 class ShrinkingTest : FunSpec() {
@@ -56,5 +59,23 @@ Shrink result (after 14 shrinks) => "<<<<"
             """.trim()
          )
       }
+
+      test("withEdgecases should maintain the shrinker") {
+         val arb = Arb.int().withEdgecases(5, 6)
+
+         val stdout = captureStandardOut {
+            shouldThrowAny {
+               checkAll(PropTestConfig(seed = 324236), arb, arb) { a, b ->
+                  (a + b) shouldBeLessThan 4
+               }
+            }
+         }
+
+         stdout.shouldContain(
+            """Attempting to shrink arg -245346456
+Shrink #1: 0 fail
+Shrink result (after 1 shrinks) => 0"""
+         )
+      }
    }
 }