feat(ChunkWriteService): support iterable directly

wrk-flow · Oct 17, 2024 · 9e21547 · 9e21547
1 parent 3f4f982
commit 9e21547
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 6 deletions.
diff --git a/src/Database/Contracts/ChunkWriteServiceContract.php b/src/Database/Contracts/ChunkWriteServiceContract.php
@@ -5,14 +5,16 @@
 namespace LaraStrict\Database\Contracts;
 
 use Closure;
-use Generator;
 use Illuminate\Database\Eloquent\Model;
 use LaraStrict\Database\Entities\ChunkWriteStateEntity;
 
+/**
+ * @phpstan-type TData iterable<int|string, Model>
+ */
 interface ChunkWriteServiceContract
 {
     /**
-     * @param Closure(): Generator<int, Model> $closure
+     * @param Closure(): TData|TData $closure
      */
-    public function write(Closure $closure): ChunkWriteStateEntity;
+    public function write(Closure|iterable $closure): ChunkWriteStateEntity;
 }
diff --git a/src/Database/Services/ChunkWriteService.php b/src/Database/Services/ChunkWriteService.php
@@ -12,11 +12,12 @@
 
 final class ChunkWriteService implements ChunkWriteServiceContract
 {
-    public function write(Closure $closure, int $batchSize = 0): ChunkWriteStateEntity
+    public function write(Closure|iterable $closure, int $batchSize = 0): ChunkWriteStateEntity
     {
         $writeState = new ChunkWriteStateEntity(batchSize: $batchSize);
 
-        foreach ($closure() as $model) {
+        $source = $closure instanceof Closure ? $closure() : $closure;
+        foreach ($source as $model) {
             $this->add($model, $writeState);
         }
 

diff --git a/tests/Unit/Database/Services/ChunkWriteServiceTest.php b/tests/Unit/Database/Services/ChunkWriteServiceTest.php
@@ -5,12 +5,16 @@
 namespace Tests\LaraStrict\Unit\Database\Services;
 
 use Closure;
+use LaraStrict\Database\Contracts\ChunkWriteServiceContract;
 use LaraStrict\Database\Entities\ChunkWriteStateEntity;
 use LaraStrict\Database\Services\ChunkWriteService;
 use LaraStrict\Tests\Traits\SqlTestEnable;
 use PHPUnit\Framework\Assert;
 use PHPUnit\Framework\TestCase;
 
+/**
+ * @phpstan-import-type TData from ChunkWriteServiceContract
+ */
 final class ChunkWriteServiceTest extends TestCase
 {
     use SqlTestEnable;
@@ -38,6 +42,14 @@ static function () {
                     );
                 },
             ],
+            [
+                static function (self $self) {
+                    $self->assert(
+                        new ChunkWriteStateEntity(32768, TestModel::class, 3, 2),
+                        [new TestModel(), new TestModel(), new TestModel()],
+                    );
+                },
+            ],
         ];
     }
 
@@ -50,7 +62,10 @@ public function test(Closure $assert): void
         $assert($this);
     }
 
-    public function assert(ChunkWriteStateEntity $expected, Closure $data): void
+    /**
+     * @param Closure(): TData|TData $data
+     */
+    public function assert(ChunkWriteStateEntity $expected, Closure|iterable $data): void
     {
         $state = (new ChunkWriteService())->write($data);
         Assert::assertEquals($expected, $state);