diff --git a/core/modules/migrate_drupal_ui/tests/src/Functional/MigrateUpgradeTestBase.php b/core/modules/migrate_drupal_ui/tests/src/Functional/MigrateUpgradeTestBase.php
index 73158feef0..bc73aba4c4 100644
--- a/core/modules/migrate_drupal_ui/tests/src/Functional/MigrateUpgradeTestBase.php
+++ b/core/modules/migrate_drupal_ui/tests/src/Functional/MigrateUpgradeTestBase.php
@@ -232,9 +232,18 @@ protected function assertMigrationResults(array $expected_counts, $version) {
     // Have to reset all the statics after migration to ensure entities are
     // loadable.
     $this->resetAll();
-    foreach (array_keys(\Drupal::entityTypeManager()->getDefinitions()) as $entity_type) {
+    // Check that the expected number of entities is the same as the actual
+    // number of entities.
+    $entity_definitions = array_keys(\Drupal::entityTypeManager()->getDefinitions());
+    $expected_count_keys = array_keys($expected_counts);
+    sort($entity_definitions);
+    sort($expected_count_keys);
+    $this->assertSame($expected_count_keys, $entity_definitions);
+
+    // Assert the correct number of entities exist.
+    foreach ($entity_definitions as $entity_type) {
       $real_count = (int) \Drupal::entityQuery($entity_type)->count()->execute();
-      $expected_count = isset($expected_counts[$entity_type]) ? $expected_counts[$entity_type] : 0;
+      $expected_count = $expected_counts[$entity_type];
       $this->assertSame($expected_count, $real_count, "Found $real_count $entity_type entities, expected $expected_count.");
     }
 
diff --git a/core/modules/migrate_drupal_ui/tests/src/Functional/d6/MigrateUpgrade6Test.php b/core/modules/migrate_drupal_ui/tests/src/Functional/d6/MigrateUpgrade6Test.php
index 29219d8ff7..75cff1ca17 100644
--- a/core/modules/migrate_drupal_ui/tests/src/Functional/d6/MigrateUpgrade6Test.php
+++ b/core/modules/migrate_drupal_ui/tests/src/Functional/d6/MigrateUpgrade6Test.php
@@ -65,6 +65,7 @@ protected function getEntityCounts() {
       // migration creates 12 comment types, one per node type.
       'comment_type' => 13,
       'contact_form' => 5,
+      'contact_message' => 0,
       'configurable_language' => 5,
       'editor' => 2,
       'field_config' => 90,
@@ -73,7 +74,6 @@ protected function getEntityCounts() {
       'filter_format' => 7,
       'image_style' => 5,
       'language_content_settings' => 10,
-      'migration' => 105,
       'node' => 18,
       // The 'book' module provides the 'book' node type, and the migration
       // creates 12 node types.
diff --git a/core/modules/migrate_drupal_ui/tests/src/Functional/d7/MigrateUpgrade7Test.php b/core/modules/migrate_drupal_ui/tests/src/Functional/d7/MigrateUpgrade7Test.php
index b1718011f0..1579c830ba 100644
--- a/core/modules/migrate_drupal_ui/tests/src/Functional/d7/MigrateUpgrade7Test.php
+++ b/core/modules/migrate_drupal_ui/tests/src/Functional/d7/MigrateUpgrade7Test.php
@@ -70,6 +70,7 @@ protected function getEntityCounts() {
       // and 'fr'.
       'configurable_language' => 5,
       'contact_form' => 3,
+      'contact_message' => 0,
       'editor' => 2,
       'field_config' => 67,
       'field_storage_config' => 50,
@@ -77,7 +78,6 @@ protected function getEntityCounts() {
       'filter_format' => 7,
       'image_style' => 6,
       'language_content_settings' => 6,
-      'migration' => 73,
       'node' => 5,
       'node_type' => 6,
       'rdf_mapping' => 8,