Adds endpoint setting to prefetchiterator plugin (#212)

library/Solarium/Plugin/PrefetchIterator.php

@@ -43,6 +43,7 @@ use Solarium\Core\Plugin\Plugin;
 use Solarium\QueryType\Select\Query\Query as SelectQuery;
 use Solarium\QueryType\Select\Result\Result as SelectResult;
 use Solarium\QueryType\Select\Result\DocumentInterface;
+use Solarium\Core\Client\Endpoint;
 
 /**
  * Prefetch plugin
@@ -142,6 +143,29 @@ class PrefetchIterator extends Plugin implements \Iterator, \Countable
         return $this->query;
     }
 
+    /**
+     * Set endpoint to use
+     *
+     * This overwrites any existing endpoint
+     *
+     * @param string|Endpoint $endpoint
+     */
+    public function setEndpoint($endpoint)
+    {
+        return $this->setOption('endpoint', $endpoint);
+    }
+
+    /**
+     * Get endpoint setting
+     *
+     * @return string|Endpoint|null
+     */
+    public function getEndpoint()
+    {
+        return $this->getOption('endpoint');
+    }
+
+
     /**
      * Countable implementation
      *
@@ -223,7 +247,7 @@ class PrefetchIterator extends Plugin implements \Iterator, \Countable
     protected function fetchNext()
     {
         $this->query->setStart($this->start)->setRows($this->getPrefetch());
-        $this->result = $this->client->execute($this->query);
+        $this->result = $this->client->execute($this->query, $this->getOption('endpoint'));
         $this->documents = $this->result->getDocuments();
         $this->start += $this->getPrefetch();
     }

tests/Solarium/Tests/Plugin/PrefetchIteratorTest.php

@@ -79,7 +79,10 @@ class PrefetchIteratorTest extends \PHPUnit_Framework_TestCase
     {
         $result = $this->getResult();
         $mockClient = $this->getMock('Solarium\Core\Client\Client', array('execute'));
-        $mockClient->expects($this->exactly(1))->method('execute')->will($this->returnValue($result));
+        $mockClient->expects($this->exactly(1))
+                   ->method('execute')
+                   ->with($this->equalTo($this->query), $this->equalTo(null))
+                   ->will($this->returnValue($result));
 
         $this->plugin->initPlugin($mockClient, array());
         $this->plugin->setQuery($this->query);
@@ -178,6 +181,41 @@ class PrefetchIteratorTest extends \PHPUnit_Framework_TestCase
 
         return new SelectDummy(1, 12, $numFound, $docs, array());
     }
+
+    public function testSetAndGetEndpointAsString()
+    {
+        $this->assertEquals(null, $this->plugin->getEndpoint());
+        $this->plugin->setEndpoint('s1');
+        $this->assertEquals('s1', $this->plugin->getEndpoint());
+    }
+
+    public function testWithSpecificEndpoint()
+    {
+        $result = $this->getResult();
+        $mockClient = $this->getMock('Solarium\Core\Client\Client', array('execute'));
+        $mockClient->expects($this->exactly(1))
+                   ->method('execute')
+                   ->with($this->equalTo($this->query), $this->equalTo('s2'))
+                   ->will($this->returnValue($result));
+
+        $this->plugin->initPlugin($mockClient, array());
+        $this->plugin->setQuery($this->query)->setEndpoint('s2');
+        $this->assertEquals(5, count($this->plugin));
+    }
+
+    public function testWithSpecificEndpointOption()
+    {
+        $result = $this->getResult();
+        $mockClient = $this->getMock('Solarium\Core\Client\Client', array('execute'));
+        $mockClient->expects($this->exactly(1))
+                   ->method('execute')
+                   ->with($this->equalTo($this->query), $this->equalTo('s3'))
+                   ->will($this->returnValue($result));
+
+        $this->plugin->initPlugin($mockClient, array('endpoint' => 's3'));
+        $this->plugin->setQuery($this->query);
+        $this->assertEquals(5, count($this->plugin));
+    }
 }
 
 class SelectDummy extends Result