remove callOnly in favor of callingOnly with callback param

Author: d8vjork

src/UnionBuilder.php

@@ -2,6 +2,7 @@
 
 namespace OpenSoutheners\LaravelEloquentUnionBuilder;
 
+use Closure;
 use Exception;
 use Illuminate\Database\Eloquent\Builder;
 use Illuminate\Database\SQLiteConnection;
@@ -28,11 +29,6 @@ final class UnionBuilder
      */
     protected $selectModelsColumns = [];
 
-    /**
-     * @var array<string>
-     */
-    protected $callingOnly = [];
-
     /**
      * Construct new instance of class.
      *
@@ -177,11 +173,12 @@ public function getAllSelectedColumns()
     public function add(Builder $builder, array $columns = [])
     {
         $builderModel = $builder->getModel();
+        $builderModelClass = get_class($builder->getModel());
 
-        $this->selectModelsColumns[get_class($builderModel)] = $columns
+        $this->selectModelsColumns[$builderModelClass] = $columns
             ?: Schema::getColumnListing($builderModel->getTable());
 
-        $this->builders[] = $builder;
+        $this->builders[$builderModelClass] = $builder;
 
         return $this;
     }
@@ -190,11 +187,12 @@ public function add(Builder $builder, array $columns = [])
      * Forward next call only on the specified model's builder.
      *
      * @param  class-string<\Illuminate\Database\Eloquent\Model>  $model
+     * @param  \Closure  $callback
      * @return $this
      */
-    public function callOnly($model)
+    public function callingOnly(string $model, Closure $callback)
     {
-        $this->callingOnly[] = $model;
+        $callback($this->builders[$model]);
 
         return $this;
     }
@@ -209,15 +207,9 @@ public function callOnly($model)
     public function __call($method, $arguments)
     {
         foreach ($this->builders as $builder) {
-            if (! empty($this->callingOnly) && ! in_array(get_class($builder->getModel()), $this->callingOnly)) {
-                continue;
-            }
-
             $this->forwardCallTo($builder, $method, $arguments);
         }
 
-        $this->callingOnly = [];
-
         return $this;
     }
 }

tests/UnionBuilderTest.php

@@ -3,6 +3,7 @@
 namespace OpenSoutheners\LaravelEloquentUnionBuilder\Tests;
 
 use Exception;
+use Illuminate\Database\Eloquent\Builder;
 use OpenSoutheners\LaravelEloquentUnionBuilder\Tests\Fixtures\Post;
 use OpenSoutheners\LaravelEloquentUnionBuilder\Tests\Fixtures\Tag;
 use OpenSoutheners\LaravelEloquentUnionBuilder\Tests\Fixtures\User;
@@ -119,9 +120,9 @@ public function testUnionBuilderSearchThrowsExceptionWhenNonSearchableModelFound
     public function testUnionBuilderCallOnlyWhereOnPostReturnsFilteredPostsOnly()
     {
         $queryResults = UnionBuilder::from([Tag::class, Post::class])
-            ->callOnly(Post::class)
-            ->where('slug', 'hello-world')
-            ->get();
+            ->callingOnly(Post::class, function (Builder $query) {
+                $query->where('slug', 'hello-world');
+            })->get();
 
         $this->assertCount(3, $queryResults);