diff --git a/src/Controller/AbstractEntityController.php b/src/Controller/AbstractEntityController.php
index 2862809..253f798 100644
--- a/src/Controller/AbstractEntityController.php
+++ b/src/Controller/AbstractEntityController.php
@@ -59,6 +59,13 @@ public function getEntity() {
     return $this->entity;
   }
 
+  /**
+   * @param \Drupal\Core\Entity\EntityInterface $entity
+   */
+  public function setEntity(EntityInterface $entity) {
+    $this->entity = $entity;
+  }
+
   /**
    * @return string
    */
diff --git a/src/Factory/EntityControllerFactory.php b/src/Factory/EntityControllerFactory.php
index f8144ae..aa8ead0 100644
--- a/src/Factory/EntityControllerFactory.php
+++ b/src/Factory/EntityControllerFactory.php
@@ -7,14 +7,28 @@
 
 namespace Drupal\cw_tool\Factory;
 
+use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\cw_tool\Controller\AbstractEntityController;
 
+/**
+ * Class EntityControllerFactory
+ * @package Drupal\cw_tool\Factory
+ *
+ * The purpose of this class is to create entity controllers.
+ */
 class EntityControllerFactory {
 
   const MASTER_CONTROLLER_CLASS = 'Drupal\cw_tool\Controller\AbstractEntityController';
 
+  /**
+   * @var string
+   */
   private $entityType;
 
+  /**
+   * @var string
+   */
   private $controllerClass;
 
   /**
@@ -22,6 +36,18 @@ class EntityControllerFactory {
    */
   private $entityManager;
 
+  /**
+   * EntityControllerFactory constructor.
+   *
+   * @param \Drupal\Core\Entity\EntityManagerInterface $entityManager
+   *   Drupal entity manager.
+   * @param string $entityType
+   *   Entity type.
+   * @param string $controllerClass
+   *   Actual entity controller class.
+   *
+   * @throws \InvalidArgumentException
+   */
   public function __construct(EntityManagerInterface $entityManager, $entityType, $controllerClass) {
     $this->entityType = $entityType;
 
@@ -33,8 +59,32 @@ public function __construct(EntityManagerInterface $entityManager, $entityType,
     $this->entityManager = $entityManager;
   }
 
+  /**
+   * Factory method to provide the controller by id.
+   *
+   * @param int|null $id
+   *   ID of the entity.
+   *
+   * @return AbstractEntityController
+   *   Instance of the controller.
+   */
   public function initWithID($id) {
     return new $this->controllerClass($this->entityManager, $this->entityType, $id);
   }
 
+  /**
+   * Factory method to provide the controller by id.
+   *
+   * @param \Drupal\Core\Entity\EntityInterface $entity
+   *   Entity.
+   *
+   * @return AbstractEntityController
+   *   Instance of the controller.
+   */
+  public function initWithEntity(EntityInterface $entity) {
+    $controller = $this->initWithID($entity->id());
+    $controller->setEntity($entity);
+    return $controller;
+  }
+
 }
diff --git a/tests/src/Unit/Controller/AbstractEntityControllerTest.php b/tests/src/Unit/Controller/AbstractEntityControllerTest.php
index 7d326dd..16bf66c 100644
--- a/tests/src/Unit/Controller/AbstractEntityControllerTest.php
+++ b/tests/src/Unit/Controller/AbstractEntityControllerTest.php
@@ -5,7 +5,7 @@
 
 namespace Drupal\Tests\cw_tool\Unit\Controller;
 
-use Drupal\KernelTests\KernelTestBase;
+use Drupal\Tests\UnitTestCase;
 use Drupal\Tests\cw_tool\Fixtures\Controller\MinimalEntityController;
 
 /**
@@ -16,7 +16,7 @@
  *
  * @group cw_tool
  */
-class AbstractEntityControllerTest extends KernelTestBase {
+class AbstractEntityControllerTest extends UnitTestCase {
 
   /**
    * @var \PHPUnit_Framework_MockObject_MockObject
@@ -29,6 +29,7 @@ class AbstractEntityControllerTest extends KernelTestBase {
   private $entityStorageMock;
 
   public function setUp() {
+    include_once __DIR__ . '/../../Fixtures/Controller/MinimalEntityController.php';
     $this->entityManagerMock = $this->getMock('Drupal\Core\Entity\EntityManagerInterface');
     $this->entityStorageMock = $this->getMock('Drupal\Core\Entity\EntityStorageInterface');
     parent::setUp();