diff --git a/core/modules/content_moderation/src/EntityOperations.php b/core/modules/content_moderation/src/EntityOperations.php
index 8964d9f..dee2773 100644
--- a/core/modules/content_moderation/src/EntityOperations.php
+++ b/core/modules/content_moderation/src/EntityOperations.php
@@ -48,6 +48,8 @@ class EntityOperations implements ContainerInjectionInterface {
   protected $tracker;
 
   /**
+   * The entity bundle information service.
+   *
    * @var \Drupal\Core\Entity\EntityTypeBundleInfoInterface
    */
   protected $bundleInfo;
@@ -63,13 +65,15 @@ class EntityOperations implements ContainerInjectionInterface {
    *   The form builder.
    * @param \Drupal\content_moderation\RevisionTrackerInterface $tracker
    *   The revision tracker.
+   * @param \Drupal\Core\Entity\EntityTypeBundleInfoInterface $bundle_info
+   *   The entity bundle information service.
    */
-  public function __construct(ModerationInformationInterface $moderation_info, EntityTypeManagerInterface $entity_type_manager, FormBuilderInterface $form_builder, RevisionTrackerInterface $tracker, EntityTypeBundleInfoInterface $bundle_indo) {
+  public function __construct(ModerationInformationInterface $moderation_info, EntityTypeManagerInterface $entity_type_manager, FormBuilderInterface $form_builder, RevisionTrackerInterface $tracker, EntityTypeBundleInfoInterface $bundle_info) {
     $this->moderationInfo = $moderation_info;
     $this->entityTypeManager = $entity_type_manager;
     $this->formBuilder = $form_builder;
     $this->tracker = $tracker;
-    $this->bundleInfo = $bundle_indo;
+    $this->bundleInfo = $bundle_info;
   }
 
   /**
diff --git a/core/modules/content_moderation/src/ModerationInformation.php b/core/modules/content_moderation/src/ModerationInformation.php
index 7e7dba7..c109b31 100644
--- a/core/modules/content_moderation/src/ModerationInformation.php
+++ b/core/modules/content_moderation/src/ModerationInformation.php
@@ -138,11 +138,7 @@ public function isLiveRevision(ContentEntityInterface $entity) {
   }
 
   /**
-   * @param $entity_type
-   * @param $bundle
-   *
-   * @return \Drupal\workflow\WorkflowInterface
-   *   The workflow entity.
+   * {@inheritdoc}
    */
   public function getWorkFlowForEntity(ContentEntityInterface $entity) {
     $bundles = $this->bundleInfo->getBundleInfo($entity->getEntityTypeId());
diff --git a/core/modules/content_moderation/src/ModerationInformationInterface.php b/core/modules/content_moderation/src/ModerationInformationInterface.php
index eac7999..cc6df1e 100644
--- a/core/modules/content_moderation/src/ModerationInformationInterface.php
+++ b/core/modules/content_moderation/src/ModerationInformationInterface.php
@@ -127,11 +127,13 @@ public function hasForwardRevision(ContentEntityInterface $entity);
   public function isLiveRevision(ContentEntityInterface $entity);
 
   /**
+   * Gets the workflow for the given content entity.
+   *
    * @param \Drupal\Core\Entity\ContentEntityInterface $entity
    *   The content entity to get the workflow for.
    *
    * @return \Drupal\workflow\WorkflowInterface|null
-   *   The workflow entity.
+   *   The workflow entity. NULL if there is no workflow.
    */
   public function getWorkFlowForEntity(ContentEntityInterface $entity);
 
diff --git a/core/modules/content_moderation/tests/src/Unit/ModerationInformationTest.php b/core/modules/content_moderation/tests/src/Unit/ModerationInformationTest.php
index 632b9bc..b7406e0 100644
--- a/core/modules/content_moderation/tests/src/Unit/ModerationInformationTest.php
+++ b/core/modules/content_moderation/tests/src/Unit/ModerationInformationTest.php
@@ -5,10 +5,12 @@
 use Drupal\content_moderation\Entity\Handler\ModerationHandler;
 use Drupal\Core\Entity\ContentEntityInterface;
 use Drupal\Core\Entity\ContentEntityType;
+use Drupal\Core\Entity\EntityStorageInterface;
 use Drupal\Core\Entity\EntityTypeBundleInfoInterface;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Session\AccountInterface;
 use Drupal\content_moderation\ModerationInformation;
+use Drupal\workflow\WorkflowInterface;
 
 /**
  * @coversDefaultClass \Drupal\content_moderation\ModerationInformation
@@ -44,13 +46,11 @@ protected function getEntityTypeManager() {
    *   The bundle ID.
    * @param string|null $workflow
    *   The workflow ID. If nul no workflow information is added to the bundle.
-   * @param string|null $default_state
-   *   The default state for the workflow.
    *
    * @return \Drupal\Core\Entity\EntityTypeManagerInterface
    *   The mocked entity type manager.
    */
-  public function setupModerationBundleInfo($bundle, $workflow = NULL, $default_state = NULL) {
+  public function setupModerationBundleInfo($bundle, $workflow = NULL) {
     $bundle_info_array = [];
     if ($workflow) {
       $bundle_info_array['workflow'] = $workflow;
@@ -65,8 +65,8 @@ public function setupModerationBundleInfo($bundle, $workflow = NULL, $default_st
    * @dataProvider providerWorkflow
    * @covers ::isModeratedEntity
    */
-  public function testIsModeratedEntity($workflow, $default_state, $expected) {
-    $moderation_information = new ModerationInformation($this->getEntityTypeManager(), $this->setupModerationBundleInfo('test_bundle', $workflow, $default_state));
+  public function testIsModeratedEntity($workflow, $expected) {
+    $moderation_information = new ModerationInformation($this->getEntityTypeManager(), $this->setupModerationBundleInfo('test_bundle', $workflow));
 
     $entity_type = new ContentEntityType([
       'id' => 'test_entity_type',
@@ -82,16 +82,39 @@ public function testIsModeratedEntity($workflow, $default_state, $expected) {
 
   /**
    * @dataProvider providerWorkflow
+   * @covers ::getWorkFlowForEntity
+   */
+  public function testGetWorkFlowForEntity($workflow) {
+    $entity_type_manager = $this->prophesize(EntityTypeManagerInterface::class);
+    if ($workflow) {
+      $workflow_entity = $this->prophesize(WorkflowInterface::class)->reveal();
+      $workflow_storage = $this->prophesize(EntityStorageInterface::class);
+      $workflow_storage->load('workflow')->willReturn($workflow_entity)->shouldBeCalled();
+      $entity_type_manager->getStorage('workflow')->willReturn($workflow_storage->reveal());
+    }
+    else {
+      $workflow_entity = NULL;
+    }
+    $moderation_information = new ModerationInformation($entity_type_manager->reveal(), $this->setupModerationBundleInfo('test_bundle', $workflow));
+    $entity = $this->prophesize(ContentEntityInterface::class);
+    $entity->getEntityTypeId()->willReturn('test_entity_type');
+    $entity->bundle()->willReturn('test_bundle');
+
+    $this->assertEquals($workflow_entity, $moderation_information->getWorkFlowForEntity($entity->reveal()));
+  }
+
+  /**
+   * @dataProvider providerWorkflow
    * @covers ::shouldModerateEntitiesOfBundle
    */
-  public function testShouldModerateEntities($workflow, $default_state, $expected) {
+  public function testShouldModerateEntities($workflow, $expected) {
     $entity_type = new ContentEntityType([
       'id' => 'test_entity_type',
       'bundle_entity_type' => 'entity_test_bundle',
       'handlers' => ['moderation' => ModerationHandler::class],
     ]);
 
-    $moderation_information = new ModerationInformation($this->getEntityTypeManager(), $this->setupModerationBundleInfo('test_bundle', $workflow, $default_state));
+    $moderation_information = new ModerationInformation($this->getEntityTypeManager(), $this->setupModerationBundleInfo('test_bundle', $workflow));
 
     $this->assertEquals($expected, $moderation_information->shouldModerateEntitiesOfBundle($entity_type, 'test_bundle'));
   }
@@ -101,8 +124,8 @@ public function testShouldModerateEntities($workflow, $default_state, $expected)
    */
   public function providerWorkflow() {
     return [
-      [NULL, NULL, FALSE],
-      ['workflow', 'draft', TRUE],
+      [NULL, FALSE],
+      ['workflow', TRUE],
     ];
   }
 
diff --git a/core/modules/workflow/src/Entity/Workflow.php b/core/modules/workflow/src/Entity/Workflow.php
index 94ee79a..140c8db 100644
--- a/core/modules/workflow/src/Entity/Workflow.php
+++ b/core/modules/workflow/src/Entity/Workflow.php
@@ -158,9 +158,10 @@ public function getStates($state_ids = NULL) {
     if ($state_ids === NULL) {
       $state_ids = array_keys($this->states);
     }
+    /** @var \Drupal\workflow\StateInterface[] $states */
     $states = array_combine($state_ids, array_map([$this, 'getState'], $state_ids));
     if (count($states) > 1) {
-      // Use array multisort to sort states by weight and then label.
+      // Sort states by weight and then label.
       $weights = $labels = [];
       foreach ($states as $id => $state) {
         $weights[$id] = $state->weight();
@@ -283,10 +284,11 @@ public function getTransitions(array $transition_ids = NULL) {
     if ($transition_ids === NULL) {
       $transition_ids = array_keys($this->transitions);
     }
+    /** @var \Drupal\workflow\TransitionInterface[] $transitions */
     $transitions = array_combine($transition_ids, array_map([$this, 'getTransition'], $transition_ids));
     if (count($transitions) > 1) {
-      // Use array multisort to sort transitions by from state weight then to
-      // state weight and then transition label.
+      // Sort transitions by from state weights and then to state weights and
+      // then transition labels.
       $from_weights = $to_weights = $labels = [];
       foreach ($transitions as $id => $transition) {
         $from_weights[$id] = $transition->from()->weight();
diff --git a/core/modules/workflow/src/Form/WorkflowStateAddForm.php b/core/modules/workflow/src/Form/WorkflowStateAddForm.php
index 0dbb8a4..a231000 100644
--- a/core/modules/workflow/src/Form/WorkflowStateAddForm.php
+++ b/core/modules/workflow/src/Form/WorkflowStateAddForm.php
@@ -62,6 +62,7 @@ public function form(array $form, FormStateInterface $form_state) {
    *   TRUE if the workflow state exists, FALSE otherwise.
    */
   public function exists($state_id) {
+    /** @var \Drupal\workflow\WorkflowInterface $original_workflow */
     $original_workflow = \Drupal::entityTypeManager()->getStorage('workflow')->loadUnchanged($this->getEntity()->id());
     return $original_workflow->hasState($state_id);
   }