dataconnect(test): add sum partition and optional variable Arbs

firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/KotestTestutilPrinters.kt

@@ -16,6 +16,7 @@
 
 package com.google.firebase.dataconnect.testutil
 
+import com.google.firebase.dataconnect.OptionalVariable
 import com.google.protobuf.Duration as DurationProto
 import io.kotest.assertions.print.Print
 import io.kotest.assertions.print.Printed
@@ -33,6 +34,7 @@ fun registerDataConnectKotestTestutilPrinters() {
   Printers.add(Triple::class, TriplePrint)
   Printers.add(Quadruple::class, QuadruplePrint)
   Printers.add(Quintuple::class, QuintuplePrint)
+  Printers.add(OptionalVariable::class, OptionalVariablePrint)
 
   try {
     Printers.add(SignificanceResult::class, SignificanceResultPrint)
@@ -101,6 +103,19 @@ private object QuintuplePrint : Print<Quintuple<*, *, *, *, *>> {
         "${third.print().value}, ${fourth.print().value}, ${fifth.print().value})"
 }
 
+private object OptionalVariablePrint : Print<OptionalVariable<*>> {
+
+  @Suppress("OVERRIDE_DEPRECATION")
+  override fun print(a: OptionalVariable<*>): Printed = a.printString.printed()
+
+  private val OptionalVariable<*>.printString: String
+    get() =
+      when (this) {
+        OptionalVariable.Undefined -> "OptionalVariable.Undefined"
+        is OptionalVariable.Value<*> -> "OptionalVariable.Value(${value.print().value})"
+      }
+}
+
 private object SignificanceResultPrint : Print<SignificanceResult> {
 
   @Suppress("OVERRIDE_DEPRECATION")

firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/property/arbitrary/PropTestConfigExts.kt

@@ -17,8 +17,16 @@
 package com.google.firebase.dataconnect.testutil.property.arbitrary
 
 import io.kotest.common.ExperimentalKotest
+import io.kotest.property.EdgeConfig
 import io.kotest.property.PropTestConfig
 
 fun PropTestConfig.withIterations(iterations: Int): PropTestConfig {
   @OptIn(ExperimentalKotest::class) return copy(iterations = iterations)
 }
+
+fun PropTestConfig.withEdgeConfig(edgeConfig: EdgeConfig): PropTestConfig {
+  @OptIn(ExperimentalKotest::class) return copy(edgeConfig = edgeConfig)
+}
+
+fun PropTestConfig.withEdgeConfigEdgeCasesOnly(): PropTestConfig =
+  withEdgeConfig(EdgeConfig(edgecasesGenerationProbability = 1.0))

firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/property/arbitrary/SumPartitionArb.kt

@@ -0,0 +1,163 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.dataconnect.testutil.property.arbitrary
+
+import io.kotest.assertions.print.print
+import io.kotest.property.Arb
+import io.kotest.property.RandomSource
+import io.kotest.property.arbitrary.arbitrary
+import io.kotest.property.arbitrary.int
+import io.kotest.property.asSample
+
+/**
+ * A Kotest [Arb] that generates random partitions of a non-negative integer [sum] into a fixed
+ * number of non-negative parts ([summandCount]).
+ *
+ * Each generated [Sample] contains a list of integers whose size is exactly [summandCount] and
+ * whose elements sum up to exactly [sum].
+ *
+ * For example, with `sum = 10` and `summandCount = 3`, generated samples could include:
+ * - `[3, 2, 5]`
+ * - `[0, 10, 0]`
+ * - `[1, 7, 2]`
+ *
+ * @param sum The target sum that all generated summands must add up to. Must be non-negative.
+ * @param summandCount The number of elements in the generated list of summands. Must be
+ * non-negative. If [summandCount] is `0`, then [sum] must also be `0`.
+ */
+class SumPartitionArb(private val sum: Int, private val summandCount: Int) :
+  Arb<SumPartitionArb.Sample>() {
+
+  init {
+    require(sum >= 0) { "invalid sum: $sum" }
+    require(summandCount >= 0) { "invalid summandCount: $summandCount" }
+    require(summandCount > 0 || sum == 0) {
+      "invalid sum/summandCount pair: sum=$sum, summandCount=$summandCount"
+    }
+    require(sum.toLong() + summandCount - 1 <= Int.MAX_VALUE) {
+      "sum+summandCount-1 exceeds Int.MAX_VALUE: sum=$sum, summandCount=$summandCount"
+    }
+  }
+
+  private val edgeCaseZeroesCountArb: Arb<Int> = run {
+    if (sum == 0) {
+      arbitrary { throw IllegalStateException("internal error h5zagzq8g4: should never get here") }
+    } else {
+      Arb.int(1 until summandCount)
+    }
+  }
+
+  override fun edgecase(rs: RandomSource): Sample? {
+    if (summandCount < 2 || sum == 0) {
+      return null
+    }
+
+    val edgeCase = Sample.EdgeCase.entries.random(rs.random)
+    val summands: List<Int> =
+      when (edgeCase) {
+        Sample.EdgeCase.Zeroes ->
+          buildList(summandCount) {
+            val zeroesCount = edgeCaseZeroesCountArb.next(rs, edgeCaseProbability = 0.3f)
+            check(zeroesCount > 0)
+            repeat(zeroesCount) { add(0) }
+            addAll(generateSummands(rs, summandCount - zeroesCount))
+            shuffle(rs.random)
+          }
+        Sample.EdgeCase.SortedAscending -> generateSummands(rs, summandCount).sorted()
+        Sample.EdgeCase.SortedDescending -> generateSummands(rs, summandCount).sortedDescending()
+      }
+
+    return Sample(summands, edgeCase)
+  }
+
+  override fun sample(rs: RandomSource): io.kotest.property.Sample<Sample> {
+    val summands = generateSummands(rs, summandCount)
+    val sample = Sample(summands, edgeCase = null)
+    return sample.asSample()
+  }
+
+  private fun generateSummands(
+    rs: RandomSource,
+    count: Int,
+  ): List<Int> {
+    if (count == 0) {
+      return emptyList()
+    }
+    if (count == 1) {
+      return listOf(sum)
+    }
+
+    val maxPosition = sum + count - 1
+    val cuts: List<Int> =
+      if (count - 1 < sum) {
+        buildSet {
+            while (size < count - 1) {
+              add(rs.random.nextInt(maxPosition))
+            }
+          }
+          .sorted()
+      } else {
+        val nonCuts =
+          buildSet {
+              while (size < sum) {
+                add(rs.random.nextInt(maxPosition))
+              }
+            }
+            .sorted()
+        buildList(count - 1) {
+          var prev = -1
+          for (nonCut in nonCuts) {
+            for (v in (prev + 1) until nonCut) {
+              add(v)
+            }
+            prev = nonCut
+          }
+          for (v in (prev + 1) until maxPosition) {
+            add(v)
+          }
+        }
+      }
+
+    return buildList(count) {
+      var prev = -1
+      for (cut in cuts) {
+        add(cut - prev - 1)
+        prev = cut
+      }
+      add(maxPosition - prev - 1)
+    }
+  }
+
+  class Sample(
+    val summands: List<Int>,
+    val edgeCase: EdgeCase?,
+  ) {
+
+    override fun equals(other: Any?) = other is Sample && other.summands == summands
+
+    override fun hashCode() = summands.hashCode()
+
+    override fun toString() =
+      "SumPartitionArb.Sample(summands=${summands.print().value}, edgeCase=${edgeCase?.name})"
+
+    enum class EdgeCase {
+      Zeroes,
+      SortedAscending,
+      SortedDescending,
+    }
+  }
+}

firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/property/arbitrary/arbs.kt

@@ -24,6 +24,7 @@ import com.google.firebase.dataconnect.CacheSettings
 import com.google.firebase.dataconnect.ConnectorConfig
 import com.google.firebase.dataconnect.DataConnectPathSegment
 import com.google.firebase.dataconnect.DataConnectSettings
+import com.google.firebase.dataconnect.OptionalVariable
 import com.google.firebase.dataconnect.testutil.ImmediateDeferred
 import com.google.firebase.dataconnect.testutil.LoggedInInternalAuthProvider
 import com.google.firebase.dataconnect.testutil.LoggedInMultiTokenInternalAuthProvider
@@ -329,6 +330,57 @@ object DataConnectArb {
     max: com.google.protobuf.Duration? = null,
   ): Arb<com.google.protobuf.Duration> =
     Arb.proto.duration(min = min, max = max).map { it.duration }
+
+  fun <T> optionalVariable(
+    arb: Arb<T>,
+    undefinedProbability: Double,
+  ): Arb<OptionalVariable<T>> {
+    require(undefinedProbability in 0.0..1.0) {
+      "invalid undefinedProbability: ${undefinedProbability.print().value}"
+    }
+
+    return arbitrary { rs ->
+      if (rs.random.nextDouble() < undefinedProbability) {
+        OptionalVariable.Undefined
+      } else {
+        OptionalVariable.Value(arb.bind())
+      }
+    }
+  }
+
+  fun <T : Any> nullableOptionalVariable(
+    arb: Arb<T>,
+    undefinedProbability: Double,
+    nullableProbability: Double,
+  ): Arb<OptionalVariable<T?>> {
+    require(undefinedProbability in 0.0..1.0) {
+      "invalid undefinedProbability: ${undefinedProbability.print().value}"
+    }
+    require(nullableProbability in 0.0..1.0) {
+      "invalid nullableProbability: ${nullableProbability.print().value}"
+    }
+    val probabilitiesSum = undefinedProbability + nullableProbability
+    require(probabilitiesSum <= 1.0) {
+      "invalid undefinedProbability/nullableProbability pair: " +
+        "their sum must be less than or equal to 1.0, " +
+        "but their sum is ${(probabilitiesSum).print().value}, " +
+        "which is ${(probabilitiesSum - 1.0).print().value} " +
+        "greater than 1.0; " +
+        "undefinedProbability=${undefinedProbability.print().value}, " +
+        "nullableProbability=${nullableProbability.print().value}"
+    }
+
+    return arbitrary { rs ->
+      val discriminator = rs.random.nextDouble()
+      if (discriminator < undefinedProbability) {
+        OptionalVariable.Undefined
+      } else if (discriminator < probabilitiesSum) {
+        OptionalVariable.Value(null)
+      } else {
+        OptionalVariable.Value(arb.bind())
+      }
+    }
+  }
 }
 
 private class DataConnectPathArb(