Unpack nested IteratorAggregate objects for Count

This allows generators (sebastianbergmann#2149) and internal Traversables (sebastianbergmann#2642) to be properly detected before attempting to restore the Iterator's position after counting (sebastianbergmann#1125).
jmikola · May 22, 2017 · 08451fe · 08451fe
1 parent b640bda
commit 08451fe
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 4 deletions.
diff --git a/src/Framework/Constraint/Count.php b/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);
             }

diff --git a/tests/Framework/Constraint/CountTest.php b/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()

diff --git a/tests/_files/TestIteratorAggregate2.php b/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;
+    }
+}