Skip to content

Commit 8a0ecc2

Browse files
committed
Fix take not completing without another emitted value
1 parent f59dcf9 commit 8a0ecc2

File tree

2 files changed

+28
-5
lines changed

2 files changed

+28
-5
lines changed

src/Pipeline.php

+9-3
Original file line numberDiff line numberDiff line change
@@ -489,12 +489,18 @@ public function take(int $count): self
489489
return $this->flatMap(static function (mixed $value) use ($count) {
490490
static $i = 0;
491491

492-
if ($i++ < $count) {
492+
if (++$i < $count) {
493493
return [$value];
494494
}
495495

496-
/** @var T[] */
497-
return [FlatMapOperation::getStopMarker()];
496+
/** @var T $stopMarker Fake stop marker as type T. */
497+
$stopMarker = FlatMapOperation::getStopMarker();
498+
499+
if ($i === $count) {
500+
return [$value, $stopMarker];
501+
}
502+
503+
return [$stopMarker];
498504
});
499505
}
500506

test/TakeTest.php

+19-2
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
use Amp\PHPUnit\AsyncTestCase;
66
use Amp\PHPUnit\TestException;
7+
use function Amp\delay;
78

89
class TakeTest extends AsyncTestCase
910
{
@@ -15,10 +16,26 @@ public function testValuesEmitted(): void
1516
self::assertSame([1, 2], $pipeline->toArray());
1617
}
1718

19+
public function testCompleteBeforeSourceCompletes()
20+
{
21+
$count = 3;
22+
$this->setTimeout(0.1 * $count + 0.1);
23+
24+
$emitted = Pipeline::fromIterable(function () use ($count): \Generator {
25+
for ($i = 0; $i < $count; ++$i) {
26+
delay(0.1);
27+
yield $i;
28+
}
29+
delay(1);
30+
})->take($count)->toArray();
31+
32+
self::assertSame(\range(0, $count - 1), $emitted);
33+
}
34+
1835
public function testPipelineFails(): void
1936
{
20-
$exception = new TestException;
21-
$source = new Queue;
37+
$exception = new TestException();
38+
$source = new Queue();
2239

2340
$iterator = $source->pipe()->take(2)->getIterator();
2441

0 commit comments

Comments
 (0)