diff --git a/core/modules/content_moderation/tests/src/Kernel/ContentModerationStateTest.php b/core/modules/content_moderation/tests/src/Kernel/ContentModerationStateTest.php
index 2651284..2629baa 100644
--- a/core/modules/content_moderation/tests/src/Kernel/ContentModerationStateTest.php
+++ b/core/modules/content_moderation/tests/src/Kernel/ContentModerationStateTest.php
@@ -38,14 +38,6 @@ class ContentModerationStateTest extends KernelTestBase {
    * {@inheritdoc}
    */
   protected function setUp() {
-
-    $entity_type = clone \Drupal::entityTypeManager()->getDefinition('entity_test_with_bundle');
-    $keys = $entity_type->getKeys();
-    $keys['revision'] = 'revision_id';
-    $entity_type->set('entity_keys', $keys);
-    \Drupal::state()->set('entity_test_with_bundle.entity_type', $entity_type);
-
-
     parent::setUp();
 
     $this->installSchema('node', 'node_access');
diff --git a/core/modules/system/tests/modules/entity_test/entity_test.module b/core/modules/system/tests/modules/entity_test/entity_test.module
index 4004259..453bdf1 100644
--- a/core/modules/system/tests/modules/entity_test/entity_test.module
+++ b/core/modules/system/tests/modules/entity_test/entity_test.module
@@ -95,9 +95,6 @@ function entity_test_entity_type_alter(array &$entity_types) {
   // Allow entity_test_update tests to override the entity type definition.
   $entity_types['entity_test_update'] = $state->get('entity_test_update.entity_type', $entity_types['entity_test_update']);
 
-  // Allow entity_test_with_bundle tests to override the entity type definition.
-  $entity_types['entity_test_with_bundle'] = $state->get('entity_test_with_bundle.entity_type', $entity_types['entity_test_with_bundle']);
-
   // Enable the entity_test_new only when needed.
   if (!$state->get('entity_test_new')) {
     unset($entity_types['entity_test_new']);
diff --git a/core/modules/system/tests/modules/entity_test/src/Entity/EntityTestWithBundle.php b/core/modules/system/tests/modules/entity_test/src/Entity/EntityTestWithBundle.php
index c668bd3..a3ed955 100644
--- a/core/modules/system/tests/modules/entity_test/src/Entity/EntityTestWithBundle.php
+++ b/core/modules/system/tests/modules/entity_test/src/Entity/EntityTestWithBundle.php
@@ -30,6 +30,7 @@
  *   entity_keys = {
  *     "id" = "id",
  *     "uuid" = "uuid",
+ *     "revision" = "revision_id",
  *     "bundle" = "type",
  *     "label" = "name",
  *     "langcode" = "langcode",