Fix take not completing without another emitted value

amphp · Dec 23, 2023 · 8a0ecc2 · 8a0ecc2
1 parent f59dcf9
commit 8a0ecc2
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/src/Pipeline.php b/src/Pipeline.php
@@ -489,12 +489,18 @@ public function take(int $count): self
         return $this->flatMap(static function (mixed $value) use ($count) {
             static $i = 0;
 
-            if ($i++ < $count) {
+            if (++$i < $count) {
                 return [$value];
             }
 
-            /** @var T[] */
-            return [FlatMapOperation::getStopMarker()];
+            /** @var T $stopMarker Fake stop marker as type T. */
+            $stopMarker = FlatMapOperation::getStopMarker();
+
+            if ($i === $count) {
+                return [$value, $stopMarker];
+            }
+
+            return [$stopMarker];
         });
     }
 

diff --git a/test/TakeTest.php b/test/TakeTest.php
@@ -4,6 +4,7 @@
 
 use Amp\PHPUnit\AsyncTestCase;
 use Amp\PHPUnit\TestException;
+use function Amp\delay;
 
 class TakeTest extends AsyncTestCase
 {
@@ -15,10 +16,26 @@ public function testValuesEmitted(): void
         self::assertSame([1, 2], $pipeline->toArray());
     }
 
+    public function testCompleteBeforeSourceCompletes()
+    {
+        $count = 3;
+        $this->setTimeout(0.1 * $count + 0.1);
+
+        $emitted = Pipeline::fromIterable(function () use ($count): \Generator {
+            for ($i = 0; $i < $count; ++$i) {
+                delay(0.1);
+                yield $i;
+            }
+            delay(1);
+        })->take($count)->toArray();
+
+        self::assertSame(\range(0, $count - 1), $emitted);
+    }
+
     public function testPipelineFails(): void
     {
-        $exception = new TestException;
-        $source = new Queue;
+        $exception = new TestException();
+        $source = new Queue();
 
         $iterator = $source->pipe()->take(2)->getIterator();