From 96e8f4b7fc5d8828a04604ebe0ef8908751a5983 Mon Sep 17 00:00:00 2001
From: Steven Schoen <steven.schoen@reddit.com>
Date: Mon, 12 Aug 2024 15:12:32 -0600
Subject: [PATCH] Fix expecting items after earlier items

---
 .../main/kotlin/strikt/assertions/Iterable.kt | 30 +++++++++++--------
 .../iterable/AssertionOrderingConstraints.kt  |  5 ++--
 .../IterableOrderingConstraintsAssertions.kt  |  8 +++++
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/strikt-core/src/main/kotlin/strikt/assertions/Iterable.kt b/strikt-core/src/main/kotlin/strikt/assertions/Iterable.kt
index 0363daf5..664d53d0 100644
--- a/strikt-core/src/main/kotlin/strikt/assertions/Iterable.kt
+++ b/strikt-core/src/main/kotlin/strikt/assertions/Iterable.kt
@@ -1,6 +1,5 @@
 package strikt.assertions
 
-import strikt.api.Assertion
 import strikt.api.Assertion.Builder
 import strikt.internal.iterable.ElementWithOrderingConstraints
 import strikt.internal.iterable.OrderingConstraint
@@ -654,20 +653,27 @@ fun <T: Iterable<E>, E> Builder<T>.containsWithOrderingConstraints(
     allSections.flatMap { it.elementsWithConstraints.map { it.element } }
   ) {
     var elementsConsumed = 0
-    fun consumeElementsAndAssertSection(section: SectionAssertionSpec<E>): Assertion.Builder<out Iterable<E>> {
+    fun consumeElementsAndAssertSection(section: SectionAssertionSpec<E>): Builder<out Iterable<E>> {
       val sectionElementCount: Int
       val assertion = when (val endDefinedBy = section.endDefinedBy) {
-        SectionAssertionSpec.EndDefinition.DeclaredElementCount -> {
-          // Define the section by the number of elements declared with `expect`
-          sectionElementCount = section.elementsWithConstraints.size
-          get("section %s") { drop(elementsConsumed).take(sectionElementCount) }
-            .assert("has %s elements", sectionElementCount) {
-              if (it.size == sectionElementCount) {
-                pass()
-              } else {
-                fail(it.size, "only %s elements left in list")
+        SectionAssertionSpec.EndDefinition.Undefined -> {
+          if (section === allSections.last()) {
+            // Define the section by taking everything
+            val remainingElements = subject.drop(elementsConsumed)
+            sectionElementCount = remainingElements.size
+            get("section %s") { remainingElements }
+          } else {
+            // Define the section by the number of elements declared with `expect`
+            sectionElementCount = section.elementsWithConstraints.size
+            get("section %s") { drop(elementsConsumed).take(sectionElementCount) }
+              .assert("has %s elements", sectionElementCount) {
+                if (it.size == sectionElementCount) {
+                  pass()
+                } else {
+                  fail(it.size, "only %s elements left in list")
+                }
               }
-            }
+          }
         }
         is SectionAssertionSpec.EndDefinition.DeclaredElement -> {
           // Define the section by taking everything until the end element
diff --git a/strikt-core/src/main/kotlin/strikt/internal/iterable/AssertionOrderingConstraints.kt b/strikt-core/src/main/kotlin/strikt/internal/iterable/AssertionOrderingConstraints.kt
index 1db4dbb0..fd323568 100644
--- a/strikt-core/src/main/kotlin/strikt/internal/iterable/AssertionOrderingConstraints.kt
+++ b/strikt-core/src/main/kotlin/strikt/internal/iterable/AssertionOrderingConstraints.kt
@@ -29,6 +29,7 @@ internal class OrderingConstraintsAssertScopeImpl<E>: OrderingConstraintsAssertS
 
   override fun expectNoFurtherElements() {
     expectsNoFurtherElements = true
+    currentBuildingSection.expectsNoOtherElements = true
   }
 
   override fun startNewSection() = endSectionIfInProgress()
@@ -91,11 +92,11 @@ internal data class ElementWithOrderingConstraints<E>(val element: E, val constr
 
 internal class SectionAssertionSpec<E> {
   val elementsWithConstraints = mutableListOf<ElementWithOrderingConstraints<E>>()
-  var endDefinedBy: EndDefinition<E> = EndDefinition.DeclaredElementCount
+  var endDefinedBy: EndDefinition<E> = EndDefinition.Undefined
   var expectsNoOtherElements = false
 
   sealed class EndDefinition<out E> {
-    data object DeclaredElementCount : EndDefinition<Nothing>()
+    data object Undefined : EndDefinition<Nothing>()
     data class DeclaredElement<E>(val element: E) : EndDefinition<E>()
   }
 }
diff --git a/strikt-core/src/test/kotlin/strikt/assertions/IterableOrderingConstraintsAssertions.kt b/strikt-core/src/test/kotlin/strikt/assertions/IterableOrderingConstraintsAssertions.kt
index 7b8ee2f6..b0d79b8b 100644
--- a/strikt-core/src/test/kotlin/strikt/assertions/IterableOrderingConstraintsAssertions.kt
+++ b/strikt-core/src/test/kotlin/strikt/assertions/IterableOrderingConstraintsAssertions.kt
@@ -276,6 +276,14 @@ class IterableOrderingConstraintsAssertions {
     }
   }
 
+  @Test
+  fun passesWithElementBeforeExpectedElement() {
+    expectThat(listOf("a", "b"))
+      .containsWithOrderingConstraints {
+        expect("b")
+      }
+  }
+
   @Test
   fun failsWithDuplicateElement() {
     assertThrows<AssertionError> {