Unpack nested IteratorAggregate objects for Count

This allows generators (#2149) and internal Traversables (#2642) to be properly detected before attempting to restore the Iterator's position after counting (#1125).

Author: jmikola

src/Framework/Constraint/Count.php

@@ -56,12 +56,12 @@ protected function getCountOf($other)
         }
 
         if ($other instanceof Traversable) {
-            if ($other instanceof IteratorAggregate) {
-                $iterator = $other->getIterator();
-            } else {
-                $iterator = $other;
+            while ($other instanceof IteratorAggregate) {
+                $other = $other->getIterator();
             }
 
+            $iterator = $other;
+
             if ($iterator instanceof Generator) {
                 return $this->getCountOfGenerator($iterator);
             }

tests/Framework/Constraint/CountTest.php

@@ -25,9 +25,11 @@ public function testCount()
         $countConstraint = new Count(2);
         $it              = new \TestIterator([1, 2]);
         $ia              = new \TestIteratorAggregate($it);
+        $ia2             = new \TestIteratorAggregate2($ia);
 
         $this->assertTrue($countConstraint->evaluate($it, '', true));
         $this->assertTrue($countConstraint->evaluate($ia, '', true));
+        $this->assertTrue($countConstraint->evaluate($ia2, '', true));
     }
 
     public function testCountDoesNotChangeIteratorKey()
@@ -78,6 +80,23 @@ public function testCountDoesNotChangeIteratorKey()
         $it->next();
         $countConstraint->evaluate($ia, '', true);
         $this->assertFalse($it->valid());
+
+        // test with nested IteratorAggregate
+        $it = new \TestIterator([1, 2]);
+        $ia = new \TestIteratorAggregate($it);
+        $ia2 = new \TestIteratorAggregate2($ia);
+
+        $countConstraint = new Count(2);
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertEquals(1, $it->current());
+
+        $it->next();
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertEquals(2, $it->current());
+
+        $it->next();
+        $countConstraint->evaluate($ia2, '', true);
+        $this->assertFalse($it->valid());
     }
 
     public function testCountGeneratorsDoNotRewind()

tests/_files/TestIteratorAggregate2.php

@@ -0,0 +1,19 @@
+<?php
+
+/* This class is used for testing a chain of IteratorAggregate objects, since
+ * PHP does allow IteratorAggregate::getIterator() to return an instance of the
+ * same class. */
+class TestIteratorAggregate2 implements IteratorAggregate
+{
+    private $traversable;
+
+    public function __construct(\Traversable $traversable)
+    {
+        $this->traversable = $traversable;
+    }
+
+    public function getIterator()
+    {
+        return $this->traversable;
+    }
+}