diff --git a/role_delegation.module b/role_delegation.module
index 9f3ea9e..8010d38 100644
--- a/role_delegation.module
+++ b/role_delegation.module
@@ -11,13 +11,11 @@
  * without needing access to the user edit form.
  */
 
-use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Render\Element;
 use Drupal\Core\Routing\RouteMatchInterface;
-use Drupal\Core\Session\AccountInterface;
-use Drupal\role_delegation\RoleDelegationHelper;
 use Drupal\user\Entity;
 use Drupal\user\Entity\Role;
+use Drupal\Core\Entity\EntityInterface;
 
 /**
  * Implements hook_help().
@@ -34,7 +32,7 @@ function role_delegation_help($route_name, RouteMatchInterface $route_match) {
 /**
  * Implements hook_entity_delete().
  */
-function role_delegation_entity_delete(Drupal\Core\Entity\EntityInterface $entity) {
+function role_delegation_entity_delete(EntityInterface $entity) {
   if ($entity->getEntityTypeId() == 'user_role') {
     $permission = "assign {$entity->id()} role";
 
diff --git a/role_delegation.permissions.yml b/role_delegation.permissions.yml
index bec36ec..f126679 100644
--- a/role_delegation.permissions.yml
+++ b/role_delegation.permissions.yml
@@ -4,4 +4,4 @@ assign all roles:
   restrict access: TRUE
 
 permission_callbacks:
-  - \Drupal\role_delegation\RoleDelegationPermissions::rolePermissions
+  - \Drupal\role_delegation\RoleDelegation::rolePermissions
diff --git a/role_delegation.services.yml b/role_delegation.services.yml
index e0471d3..31b6838 100644
--- a/role_delegation.services.yml
+++ b/role_delegation.services.yml
@@ -1,6 +1,8 @@
 services:
   role_delegation.access_checker:
     class: Drupal\role_delegation\Access\RoleDelegationAccessCheck
-    arguments: ['@current_user']
+    arguments: ['@current_user', '@role_delegation']
     tags:
       - { name: access_check, applies_to: _role_delegation_access_check }
+  role_delegation:
+    class: Drupal\role_delegation\RoleDelegation
diff --git a/src/Access/RoleDelegationAccessCheck.php b/src/Access/RoleDelegationAccessCheck.php
index e4e9542..53186e6 100644
--- a/src/Access/RoleDelegationAccessCheck.php
+++ b/src/Access/RoleDelegationAccessCheck.php
@@ -1,4 +1,5 @@
 <?php
+
 /**
  * @file
  * Contains \Drupal\role_delegation\Access\RoleDelegationAccessCheck.
@@ -9,7 +10,7 @@ namespace Drupal\role_delegation\Access;
 use Drupal\Core\Access\AccessResult;
 use Drupal\Core\Routing\Access\AccessInterface;
 use Drupal\Core\Session\AccountInterface;
-use Drupal\role_delegation\RoleDelegationPermissions;
+use Drupal\role_delegation\RoleDelegationInterface;
 
 /**
  * Checks access for displaying configuration edit user pages.
@@ -17,6 +18,25 @@ use Drupal\role_delegation\RoleDelegationPermissions;
 class RoleDelegationAccessCheck implements AccessInterface {
 
   /**
+   * The role delegation service.
+   *
+   * @var \Drupal\role_delegation\RoleDelegationInterface
+   */
+  protected $roleDelegation;
+
+  /**
+   * The Role Delegation access check.
+   *
+   * @param \Drupal\Core\Session\AccountInterface $account
+   *   The account we're checking access for.
+   * @param \Drupal\role_delegation\RoleDelegationInterface $role_delegation
+   *   The role delegation service.
+   */
+  public function __construct(AccountInterface $account, RoleDelegationInterface $role_delegation) {
+    $this->roleDelegation = $role_delegation;
+  }
+
+  /**
    * Custom access check for the /user/%/roles page.
    *
    * @param \Drupal\Core\Session\AccountInterface $account
@@ -28,8 +48,7 @@ class RoleDelegationAccessCheck implements AccessInterface {
   public function access(AccountInterface $account) {
     // If the user has any of the "assign custom role" permissions then we give
     // them access to the form.
-    $perms = new RoleDelegationPermissions();
-    foreach ($perms->rolePermissions() as $perm => $title) {
+    foreach ($this->roleDelegation->rolePermissions() as $perm => $title) {
       if ($account->hasPermission($perm)) {
         return AccessResult::allowed()->cachePerPermissions();
       }
diff --git a/src/Form/RoleDelegationSettingsForm.php b/src/Form/RoleDelegationSettingsForm.php
index d8c8e34..de0685d 100644
--- a/src/Form/RoleDelegationSettingsForm.php
+++ b/src/Form/RoleDelegationSettingsForm.php
@@ -1,15 +1,17 @@
 <?php
+
 /**
  * @file
  * Contains \Drupal\role_delegation\Form\RoleDelegationSettingsForm.
  */
+
 namespace Drupal\role_delegation\Form;
 
 use Drupal\Core\Form\FormBase;
 use Drupal\Core\Form\FormStateInterface;
-use Drupal\role_delegation\RoleDelegationHelper;
-use Drupal\user\UserInterface;
-
+use Drupal\Core\Session\AccountInterface;
+use Drupal\role_delegation\RoleDelegationInterface;
+use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
  * Configure book settings for this site.
@@ -17,17 +19,57 @@ use Drupal\user\UserInterface;
 class RoleDelegationSettingsForm extends FormBase {
 
   /**
+   * @var \Drupal\Core\Session\AccountInterface
+   */
+  protected $currentUser;
+
+  /**
+   * @var \Drupal\role_delegation\RoleDelegationInterface
+   */
+  protected $roleDelegation;
+
+  /**
+   * The roles page setting form.
+   *
+   * @param \Drupal\Core\Session\AccountInterface $current_user
+   *   The current user viewing the form.
+   */
+  public function __construct(RoleDelegationInterface $role_delegation, AccountInterface $current_user) {
+    $this->roleDelegation = $role_delegation;
+    $this->currentUser = $current_user;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public static function create(ContainerInterface $container) {
+    return new static(
+      $container->get('role_delegation'),
+      $container->get('current_user')
+    );
+  }
+
+  /**
    * {@inheritdoc}
    */
   public function getFormId() {
-    return 'role_delegation_admin_settings';
+    return 'role_delegation_role_assign_form';
   }
 
   /**
    * {@inheritdoc}
    */
-  public function buildForm(array $form, FormStateInterface $form_state, UserInterface $user = NULL) {
-    $form = RoleDelegationHelper::addRoleDelegationElement($form, $user);
+  public function buildForm(array $form, FormStateInterface $form_state, AccountInterface $user = NULL) {
+    $current_roles = $user->getRoles(TRUE);
+    $current_roles = array_combine($current_roles, $current_roles);
+
+    $form['account']['roles_change'] = array(
+      '#type' => 'checkboxes',
+      '#title' => isset($form['account']['roles']['#title']) ? $form['account']['roles']['#title'] : t('Roles'),
+      '#options' => $this->roleDelegation->getAssignableRoles($this->currentUser),
+      '#default_value' => $current_roles,
+      '#description' => isset($form['account']['roles']['#description']) ? $form['account']['roles']['#description'] : t('Change roles assigned to user.'),
+    );
 
     $form['actions']['#type'] = 'actions';
     $form['actions']['submit'] = array(
@@ -42,22 +84,13 @@ class RoleDelegationSettingsForm extends FormBase {
    * {@inheritdoc}
    */
   public function submitForm(array &$form, FormStateInterface $form_state) {
-
-    if(!empty($form_state->getValue('account')->uid)) {
-      $account = '';
-      $uid_list = $form_state->getValue('account')->uid->getValue();
-      foreach($uid_list as $uid) {
-        $account = user_load($uid['value']);
-      }
-      foreach($form_state->getValue('roles_change') as $rid => $value) {
-        $account->addRole($rid);
-        if($value === 0) {
-          $account->removeRole($rid);
-        }
-        $account->save();
-        drupal_set_message(t('The roles have been updated.'), 'status');
-      }
+    /** @var \Drupal\user\UserInterface $account */
+    $account = $form_state->getBuildInfo()['args'][0];
+    foreach($form_state->getValue('roles_change') as $rid => $value) {
+      $value === 0 ? $account->removeRole($rid) : $account->addRole($rid);
+      $account->save();
     }
+    drupal_set_message($this->t('The roles have been updated.'), 'status');
   }
 
 }
diff --git a/src/RoleDelegation.php b/src/RoleDelegation.php
new file mode 100644
index 0000000..490b60d
--- /dev/null
+++ b/src/RoleDelegation.php
@@ -0,0 +1,53 @@
+<?php
+
+/**
+ * @file
+ * \Drupal\role_delegation\RoleDelegation.
+ */
+
+namespace Drupal\role_delegation;
+
+use Drupal\Core\Session\AccountInterface;
+use Drupal\Core\StringTranslation\StringTranslationTrait;
+
+class RoleDelegation implements RoleDelegationInterface {
+
+  use StringTranslationTrait;
+
+  /**
+   * {@inheritdoc}
+   */
+  public function rolePermissions() {
+    $perms = array();
+    foreach ($this->getAllRoles() as $rid => $role) {
+      $perms[sprintf('assign %s role', $rid)] = [
+        'title' => $this->t('Assign %role role', ['%role' => $role->label()]),
+      ];
+    }
+
+    return $perms;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getAssignableRoles(AccountInterface $account) {
+    $assignable_roles = [];
+    foreach ($this->getAllRoles() as $role) {
+      if ($account->hasPermission(sprintf('assign %s role', $role->id())) || $account->hasPermission('assign all roles')) {
+        $assignable_roles[$role->id()] = $role->label();
+      }
+    }
+    return $assignable_roles;
+  }
+
+  /**
+   * {@inheritdoc}
+   */
+  public function getAllRoles() {
+    $all_roles = user_roles(TRUE);
+    unset($all_roles[AccountInterface::AUTHENTICATED_ROLE]);
+    return $all_roles;
+  }
+
+}
diff --git a/src/RoleDelegationHelper.php b/src/RoleDelegationHelper.php
deleted file mode 100644
index 473b6df..0000000
--- a/src/RoleDelegationHelper.php
+++ /dev/null
@@ -1,69 +0,0 @@
-<?php
-/**
- * Created by PhpStorm.
- * User: ggoetlund
- * Date: 24.11.15
- * Time: 14:24
- */
-
-namespace Drupal\role_delegation;
-
-
-use Drupal\Core\Session\AccountInterface;
-use Drupal\user\UserInterface;
-
-class RoleDelegationHelper {
-
-  /**
-   * @param array $form
-   * @param \Drupal\user\UserInterface $user
-   * @return array
-   */
-  public static function addRoleDelegationElement(array $form, UserInterface $user) {
-    $current_user = \Drupal::currentUser();
-
-    $roles_current = $user->getRoles(TRUE);
-
-    $roles_delegate = array();
-
-    $roles = user_roles(TRUE);
-    unset($roles[AccountInterface::AUTHENTICATED_ROLE]);
-    unset($roles['administrator']);
-
-    foreach ($roles as $rid => $role) {
-      if ($current_user->hasPermission('assign all roles') || $current_user->hasPermission("assign {$role->get('id')} role")) {
-        $roles_delegate[$rid] = isset($form['account']['roles']['#options'][$rid]) ? $form['account']['roles']['#options'][$rid] : $role->get('id');
-      }
-    }
-
-    if (empty($roles_delegate)) {
-      // No role can be assigned.
-      return $form;
-    }
-    if (!isset($form['account'])) {
-      $form['account'] = array(
-        '#type' => 'value',
-        '#value' => $user,
-      );
-    }
-
-    $default_options = array();
-
-    foreach ($roles_current as $role) {
-      if (in_array($role, $roles_delegate)) {
-        $default_options[$role] = $role;
-      }
-    }
-
-    // Generate the form items.
-    $form['account']['roles_change'] = array(
-      '#type' => 'checkboxes',
-      '#title' => isset($form['account']['roles']['#title']) ? $form['account']['roles']['#title'] : t('Roles'),
-      '#options' => $roles_delegate,
-      '#default_value' => array_keys(array_intersect_key(array_flip($roles_current), $roles_delegate)),
-      '#description' => isset($form['account']['roles']['#description']) ? $form['account']['roles']['#description'] : t('Change roles assigned to user.'),
-    );
-
-    return $form;
-  }
-}
diff --git a/src/RoleDelegationInterface.php b/src/RoleDelegationInterface.php
new file mode 100644
index 0000000..f02b2af
--- /dev/null
+++ b/src/RoleDelegationInterface.php
@@ -0,0 +1,35 @@
+<?php
+
+namespace Drupal\role_delegation;
+
+use Drupal\Core\Session\AccountInterface;
+
+interface RoleDelegationInterface {
+
+  /**
+   * Returns an array of permissions to assign specific roles.
+   *
+   * @return array
+   *   An array of permissions in the correct format for permission_callbacks.
+   */
+  public function rolePermissions();
+
+  /**
+   * Gets the roles a user is allowed to assing.
+   *
+   * @param \Drupal\Core\Session\AccountInterface $account
+   *   The account for which you want to know which roles they can assign.
+   *
+   * @return array
+   *   An array of roles with machine names as keys and labels as values.
+   */
+  public function getAssignableRoles(AccountInterface $account);
+
+  /**
+   * Gets all roles apart from anonymous and authenticated.
+   *
+   * @return \Drupal\user\RoleInterface[]
+   *   An array of role objects.
+   */
+  public function getAllRoles();
+}
diff --git a/src/RoleDelegationPermissions.php b/src/RoleDelegationPermissions.php
deleted file mode 100644
index 18c936b..0000000
--- a/src/RoleDelegationPermissions.php
+++ /dev/null
@@ -1,43 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\role_delegation\RoleDelegationPermissions.
- */
-
-namespace Drupal\role_delegation;
-
-use Drupal\Core\Routing\UrlGeneratorTrait;
-use Drupal\Core\Session\AccountInterface;
-use Drupal\Core\StringTranslation\StringTranslationTrait;
-use Drupal\user\Entity;
-
-/**
- * Defines a class containing permission callbacks.
- */
-class RoleDelegationPermissions {
-
-  use StringTranslationTrait;
-  use UrlGeneratorTrait;
-
-  /**
-   * Returns an array of node type permissions.
-   *
-   * @return array
-   */
-  public function rolePermissions() {
-    $perms = array();
-
-    $roles = user_roles(TRUE);
-    unset($roles[AccountInterface::AUTHENTICATED_ROLE]);
-    unset($roles['administrator']);
-
-    foreach ($roles as $rid => $role) {
-      $perms["assign {$role->get('id')} role"] = array(
-        'title' => $this->t('Assign %role role', array('%role' => $role->label())),
-      );
-    }
-
-    return $perms;
-  }
-}
diff --git a/tests/Kernel/AccessTest.php b/tests/Kernel/AccessTest.php
index 68f5488..fef0ce7 100644
--- a/tests/Kernel/AccessTest.php
+++ b/tests/Kernel/AccessTest.php
@@ -11,7 +11,7 @@ use Drupal\KernelTests\KernelTestBase;
 use Drupal\simpletest\UserCreationTrait;
 
 /**
- * The AccessTest class.
+ * @coversDefaultClass \Drupal\role_delegation\Access\RoleDelegationAccessCheck
  */
 class AccessTest extends KernelTestBase {
 
@@ -27,7 +27,7 @@ class AccessTest extends KernelTestBase {
   /**
    * The Role Delegation access checker.
    *
-   * @var \Drupal\Core\Routing\Access\AccessInterface
+   * @var \Drupal\role_delegation\Access\RoleDelegationAccessCheck
    */
   protected $accessChecker;
 
@@ -46,7 +46,9 @@ class AccessTest extends KernelTestBase {
   }
 
   /**
-   * Test the access checker for user/%/roles
+   * Test the access checker for user/%/roles.
+   *
+   * @covers ::access
    */
   public function testRoleDelegationAccess() {
     // Anonymous users can never access the roles page.