From ecab7e46428cfac8af8ef315dc366caf96f24d37 Mon Sep 17 00:00:00 2001
From: Daniele Alessandri <suppakilla@gmail.com>
Date: Wed, 1 Jun 2016 12:33:20 +0200
Subject: [PATCH] Implement IteratorAggregate interface for Client.

Now it is possible to iterate over traversable aggregate connections
and get a key/value pair of $connectionId => $clientInstance for each
node.
---
 CHANGELOG.md                |  4 ++
 src/Client.php              | 21 ++++++++++-
 tests/Predis/ClientTest.php | 75 +++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ed4cf208..dbfaadad 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -39,6 +39,10 @@ v1.1.0 (2016-0x-xx)
   use the `tls` or `rediss` schemes in connection parameters along with specific
   options via the `ssl` parameter (see http://php.net/manual/context.ssl.php).
 
+- Implemented the `IteratorAggregate` interface for `Predis\Client` so now it is
+  possible to iterate over traversable aggregate connections and get a key/value
+  pair consisting of $connectionID => $clientInstance for each node.
+
 - Iterating over `Predis\Connection\Aggregate\RedisCluster` now returns all the
   connections currently mapped in the slots map instead of just the connections
   initialized in the pool. When the slots map is retrieved from Redis (which is
diff --git a/src/Client.php b/src/Client.php
index 8fd75b52..0a6c4647 100644
--- a/src/Client.php
+++ b/src/Client.php
@@ -38,7 +38,7 @@ use Predis\Transaction\MultiExec as MultiExecTransaction;
  *
  * @author Daniele Alessandri <suppakilla@gmail.com>
  */
-class Client implements ClientInterface
+class Client implements ClientInterface, \IteratorAggregate
 {
     const VERSION = '1.1.0-dev';
 
@@ -525,4 +525,23 @@ class Client implements ClientInterface
     {
         return new MonitorConsumer($this);
     }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getIterator()
+    {
+        $clients = array();
+        $connection = $this->getConnection();
+
+        if (!$connection instanceof \Traversable) {
+            throw new ClientException('The underlying connection is not traversable');
+        }
+
+        foreach ($connection as $node) {
+            $clients[(string) $node] = new static($node, $this->getOptions());
+        }
+
+        return new \ArrayIterator($clients);
+    }
 }
diff --git a/tests/Predis/ClientTest.php b/tests/Predis/ClientTest.php
index 58c226b5..c08d1cca 100644
--- a/tests/Predis/ClientTest.php
+++ b/tests/Predis/ClientTest.php
@@ -815,6 +815,55 @@ class ClientTest extends PredisTestCase
         $this->assertTrue($client->executeCommand($command));
     }
 
+    /**
+     * @group disconnected
+     */
+    public function testGetIteratorWithTraversableConnections()
+    {
+        $connection1 = $this->getMockConnection('tcp://127.0.0.1:6381');
+        $connection2 = $this->getMockConnection('tcp://127.0.0.1:6382');
+        $connection3 = $this->getMockConnection('tcp://127.0.0.1:6383');
+
+        $aggregate = new \Predis\Connection\Aggregate\PredisCluster();
+
+        $aggregate->add($connection1);
+        $aggregate->add($connection2);
+        $aggregate->add($connection3);
+
+        $client = new Client($aggregate);
+
+        $iterator = $client->getIterator();
+
+        $this->assertInstanceOf('\Predis\Client', $nodeClient = $iterator->current());
+        $this->assertSame($connection1, $nodeClient->getConnection());
+        $this->assertSame('127.0.0.1:6381', $iterator->key());
+
+        $iterator->next();
+
+        $this->assertInstanceOf('\Predis\Client', $nodeClient = $iterator->current());
+        $this->assertSame($connection2, $nodeClient->getConnection());
+        $this->assertSame('127.0.0.1:6382', $iterator->key());
+
+        $iterator->next();
+
+        $this->assertInstanceOf('\Predis\Client', $nodeClient = $iterator->current());
+        $this->assertSame($connection3, $nodeClient->getConnection());
+        $this->assertSame('127.0.0.1:6383', $iterator->key());
+    }
+
+    /**
+     * @group disconnected
+     * @expectedException \Predis\ClientException
+     * @expectedExceptionMessage The underlying connection is not traversable
+     */
+    public function testGetIteratorWithNonTraversableConnectionThrowsException()
+    {
+        $connection = $this->getMock('Predis\Connection\NodeConnectionInterface');
+        $client = new Client($connection);
+
+        $client->getIterator();
+    }
+
     // ******************************************************************** //
     // ---- HELPER METHODS ------------------------------------------------ //
     // ******************************************************************** //
@@ -843,4 +892,30 @@ class ClientTest extends PredisTestCase
 
         return $uriString;
     }
+
+    /**
+     * Returns a base mocked connection from Predis\Connection\NodeConnectionInterface.
+     *
+     * @param mixed $parameters Optional parameters.
+     *
+     * @return mixed
+     */
+    protected function getMockConnection($parameters = null)
+    {
+        $connection = $this->getMock('Predis\Connection\NodeConnectionInterface');
+
+        if ($parameters) {
+            $parameters = \Predis\Connection\Parameters::create($parameters);
+            $hash = "{$parameters->host}:{$parameters->port}";
+
+            $connection->expects($this->any())
+                       ->method('getParameters')
+                       ->will($this->returnValue($parameters));
+            $connection->expects($this->any())
+                       ->method('__toString')
+                       ->will($this->returnValue($hash));
+        }
+
+        return $connection;
+    }
 }