diff --git a/core/modules/migrate/tests/src/Unit/SqlBaseTest.php b/core/modules/migrate/tests/src/Unit/SqlBaseTest.php
index 42a03cb..57f36fd 100644
--- a/core/modules/migrate/tests/src/Unit/SqlBaseTest.php
+++ b/core/modules/migrate/tests/src/Unit/SqlBaseTest.php
@@ -56,7 +56,7 @@ protected function getSqlBase($id_map_is_sql, $with_id_map, $source_options = []
     $source_connection = $this->getMockBuilder('Drupal\Core\Database\Connection')
       ->disableOriginalConstructor()
       ->getMock();
-    $source_connection->expects($this->any())
+    $source_connection->expects($id_map_is_sql && $with_id_map ? $this->once() : $this->never())
       ->method('getConnectionOptions')
       ->willReturn($source_options);
 
@@ -64,7 +64,7 @@ protected function getSqlBase($id_map_is_sql, $with_id_map, $source_options = []
     $idmap_connection = $this->getMockBuilder('Drupal\Core\Database\Connection')
       ->disableOriginalConstructor()
       ->getMock();
-    $idmap_connection->expects($this->any())
+    $idmap_connection->expects($id_map_is_sql && $with_id_map ? $this->once() : $this->never())
       ->method('getConnectionOptions')
       ->willReturn($idmap_options);
 
@@ -72,13 +72,13 @@ protected function getSqlBase($id_map_is_sql, $with_id_map, $source_options = []
     $sql = $this->getMockBuilder('Drupal\migrate\Plugin\migrate\id_map\Sql')
       ->disableOriginalConstructor()
       ->getMock();
-    $sql->expects($this->any())
+    $sql->expects($id_map_is_sql && $with_id_map ? $this->once() : $this->never())
       ->method('getDatabase')
       ->willReturn($idmap_connection);
 
     // Setup a migration entity.
     $migration = $this->getMock('Drupal\migrate\Entity\MigrationInterface');
-    $migration->expects($this->any())
+    $migration->expects($with_id_map ? $this->once() : $this->never())
       ->method('getIdMap')
       ->willReturn($id_map_is_sql ? $sql : NULL);