diff --git a/core/modules/migrate/tests/src/MigrateExecutableTest.php b/core/modules/migrate/tests/src/MigrateExecutableTest.php
index 7b98587..6dcdea6 100644
--- a/core/modules/migrate/tests/src/MigrateExecutableTest.php
+++ b/core/modules/migrate/tests/src/MigrateExecutableTest.php
@@ -359,7 +359,7 @@ public function testImportWithRowLimit() {
     $this->executable->setRowLimit(1);
     $this->assertSame(1, $this->executable->getRowLimit());
 
-    $source = $this->getMockSource();
+    $source = $this->getMockSource(2);
 
     $row = $this->getMockBuilder('Drupal\migrate\Row')
       ->disableOriginalConstructor()
@@ -374,7 +374,7 @@ public function testImportWithRowLimit() {
       ->with(array('id' => 'test'))
       ->will($this->returnValue(array('test')));
 
-    $source->expects($this->once())
+    $source->expects($this->any())
       ->method('current')
       ->will($this->returnValue($row));
 
@@ -400,12 +400,6 @@ public function testImportWithRowLimit() {
     $this->assertSame(1, $this->executable->getTotalSuccesses());
     $this->assertSame(1, $this->executable->getTotalProcessed());
     $this->assertSame(1, $this->executable->getProcessedSinceFeedback());
-
-    $this->assertTrue(FALSE, 'this will fail. see @todo.');
-    // @todo: determine if this is the best place to test, how to
-    // how to mock multiple rows, and how to access
-    // MigrateIdMapInterface::importedCount()
-
   }
 
   /**
@@ -568,10 +562,18 @@ public function testProcessRowEmptyPipeline() {
   /**
    * Returns a mock migration source instance.
    *
+   * @param int $valid_rows
+   *   (optional) The number of valid rows the source will contain.
+   *
    * @return \Drupal\migrate\Source|\PHPUnit_Framework_MockObject_MockObject
    */
-  protected function getMockSource() {
+  protected function getMockSource($valid_rows = 1) {
     $iterator = $this->getMock('\Iterator');
+    $return_values = array();
+    for ($i = 0; $i < $valid_rows; $i++) {
+      $return_values[] = TRUE;
+    }
+    $return_values[] = FALSE;
 
     $source = $this->getMockBuilder('Drupal\migrate\Source')
       ->disableOriginalConstructor()
@@ -584,7 +586,7 @@ protected function getMockSource() {
       ->will($this->returnValue(TRUE));
     $source->expects($this->any())
       ->method('valid')
-      ->will($this->onConsecutiveCalls(TRUE, FALSE));
+      ->will(new \PHPUnit_Framework_MockObject_Stub_ConsecutiveCalls($return_values));
 
     return $source;
   }