diff --git a/core/modules/migrate/src/Plugin/migrate/source/SourcePluginBase.php b/core/modules/migrate/src/Plugin/migrate/source/SourcePluginBase.php
index 6eebdea..e6abcd5 100644
--- a/core/modules/migrate/src/Plugin/migrate/source/SourcePluginBase.php
+++ b/core/modules/migrate/src/Plugin/migrate/source/SourcePluginBase.php
@@ -174,8 +174,8 @@ public function __construct(array $configuration, $plugin_id, $plugin_definition
   /**
    * Initializes the iterator with the source data.
    *
-   * @return array
-   *   An array of the data for this source.
+   * @return \Traversable
+   *   A traversable object or array of the data for this source.
    */
   protected abstract function initializeIterator();
 
diff --git a/core/modules/migrate/src/Plugin/migrate/source/SqlBase.php b/core/modules/migrate/src/Plugin/migrate/source/SqlBase.php
index 0eac141..7b72109 100644
--- a/core/modules/migrate/src/Plugin/migrate/source/SqlBase.php
+++ b/core/modules/migrate/src/Plugin/migrate/source/SqlBase.php
@@ -232,7 +232,7 @@ protected function initializeIterator() {
   /**
    * {@inheritdoc}
    */
-  public function count() {
+  protected function doCount() {
     return $this->query()->countQuery()->execute()->fetchField();
   }
 
diff --git a/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php b/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
index 962a75d..346abfd 100644
--- a/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
+++ b/core/modules/migrate/tests/src/Unit/MigrateSqlSourceTestCase.php
@@ -2,6 +2,7 @@
 
 namespace Drupal\Tests\migrate\Unit;
 
+use Drupal\Core\Cache\CacheBackendInterface;
 use Drupal\Core\Database\Query\SelectInterface;
 use Drupal\Core\DependencyInjection\ContainerBuilder;
 use Drupal\Core\DependencyInjection\ContainerNotInitializedException;
@@ -131,6 +132,9 @@ protected function setUp() {
       ->will($this->returnValue($plugin));
     $this->source = $plugin;
     $this->expectedCount = count($this->expectedResults);
+
+    $cache = $this->getMock(CacheBackendInterface::class);
+    \Drupal::getContainer()->set('cache.migrate', $cache);
   }
 
   /**