Fix mapIds method and the corresponding test

mapIds should return Collection
mapIds needs to know the model's key name
mapIds have to have a doc block

Author: pmatseykanets

src/PostgresEngine.php

@@ -31,6 +31,11 @@ class PostgresEngine extends Engine
      */
     protected $config = [];
 
+    /**
+     * @var \Illuminate\Database\Eloquent\Model
+     */
+    protected $model;
+
     /**
      * Create a new instance of PostgresEngine.
      *
@@ -195,6 +200,11 @@ public function getTotalCount($results)
      */
     protected function performSearch(Builder $builder, $perPage = 0, $page = 1)
     {
+        // We have to preserve the model in order to allow for
+        // correct behavior of mapIds() method which currently
+        // does not revceive a model instance
+        $this->preserveModel($builder->model);
+
         $indexColumn = $this->getIndexColumn($builder->model);
 
         // Build the query
@@ -224,9 +234,19 @@ protected function performSearch(Builder $builder, $perPage = 0, $page = 1)
             ->select($query->toSql(), $bindings->all());
     }
 
+    /**
+     * Pluck and return the primary keys of the given results.
+     *
+     * @param mixed $results
+     * @return \Illuminate\Support\Collection
+     */
     public function mapIds($results)
     {
-        return collect($results)->pluck('id')->values()->all();
+        $keyName = $this->model ? $this->model->getKeyName() : 'id';
+
+        return collect($results)
+            ->pluck($keyName)
+            ->values();
     }
 
     /**
@@ -242,11 +262,11 @@ public function map($results, $model)
             return Collection::make();
         }
 
-        $results = collect($results);
-        
         $keys = $this->mapIds($results);
 
-        $models = $model->whereIn($model->getKeyName(), $keys)
+        $results = collect($results);
+
+        $models = $model->whereIn($model->getKeyName(), $keys->all())
             ->get()
             ->keyBy($model->getKeyName());
 
@@ -428,4 +448,12 @@ protected function config($key, $default = null)
     {
         return array_get($this->config, $key, $default);
     }
+
+    /**
+     * @param \Illuminate\Database\Eloquent\Model $model
+     */
+    protected function preserveModel(Model $model)
+    {
+        $this->model = $model;
+    }
 }

tests/PostgresEngineTest.php

@@ -84,23 +84,11 @@ public function test_search()
     {
         list($engine, $db) = $this->getEngine();
 
-        $db->shouldReceive('table')
-            ->andReturn($table = Mockery::mock('stdClass'));
-        $db->shouldReceive('raw')
-            ->with('plainto_tsquery(?) query')
-            ->andReturn('plainto_tsquery(?) query');
+        $table = $this->setDbExpectations($db);
 
-        $table->shouldReceive('crossJoin')->with('plainto_tsquery(?) query')->andReturnSelf()
-            ->shouldReceive('select')->with('id')->andReturnSelf()
-            ->shouldReceive('selectRaw')->with('ts_rank(searchable,query) AS rank')->andReturnSelf()
-            ->shouldReceive('selectRaw')->with('COUNT(*) OVER () AS total_count')->andReturnSelf()
-            ->shouldReceive('whereRaw')->andReturnSelf()
-            ->shouldReceive('orderBy')->with('rank', 'desc')->andReturnSelf()
-            ->shouldReceive('orderBy')->with('id')->andReturnSelf()
-            ->shouldReceive('skip')->with(0)->andReturnSelf()
+        $table->shouldReceive('skip')->with(0)->andReturnSelf()
             ->shouldReceive('limit')->with(5)->andReturnSelf()
-            ->shouldReceive('where')->with('bar', 1)->andReturnSelf()
-            ->shouldReceive('toSql');
+            ->shouldReceive('where')->with('bar', 1);
 
         $db->shouldReceive('select')->with(null, ['foo', 1]);
 
@@ -138,9 +126,19 @@ public function test_it_returns_total_count()
 
     public function test_map_ids_returns_right_key()
     {
-        list($engine) = $this->getEngine();
-        $results = [new TestModel()];
-        $this->assertEquals([1], $engine->mapIds($results));
+        list($engine, $db) = $this->getEngine();
+
+        $this->setDbExpectations($db);
+
+        $db->shouldReceive('select')
+            ->andReturn(json_decode('[{"id": 1}, {"id": 2}]'));
+
+        $builder = new Builder(new TestModel, 'foo');
+        $results = $engine->search($builder);
+        $ids = $engine->mapIds($results);
+
+        $this->assertInstanceOf(\Illuminate\Support\Collection::class, $ids);
+        $this->assertEquals([1, 2], $ids->all());
     }
 
     protected function getEngine($config = [])
@@ -153,6 +151,26 @@ protected function getEngine($config = [])
 
         return [new PostgresEngine($resolver, $config), $db];
     }
+
+    protected function setDbExpectations($db, $skip = 0, $limit = 5)
+    {
+        $db->shouldReceive('table')
+            ->andReturn($table = Mockery::mock('stdClass'));
+        $db->shouldReceive('raw')
+            ->with('plainto_tsquery(?) query')
+            ->andReturn('plainto_tsquery(?) query');
+
+        $table->shouldReceive('crossJoin')->with('plainto_tsquery(?) query')->andReturnSelf()
+            ->shouldReceive('select')->with('id')->andReturnSelf()
+            ->shouldReceive('selectRaw')->with('ts_rank(searchable,query) AS rank')->andReturnSelf()
+            ->shouldReceive('selectRaw')->with('COUNT(*) OVER () AS total_count')->andReturnSelf()
+            ->shouldReceive('whereRaw')->andReturnSelf()
+            ->shouldReceive('orderBy')->with('rank', 'desc')->andReturnSelf()
+            ->shouldReceive('orderBy')->with('id')->andReturnSelf()
+            ->shouldReceive('toSql');
+
+        return $table;
+    }
 }
 
 class TestModel extends Model