diff --git a/role_delegation.module b/role_delegation.module
index 3a24b6a..f4288e9 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\user\RoleInterface;
 
 /**
  * Implements hook_help().
@@ -25,7 +23,6 @@ use Drupal\user\Entity\Role;
 function role_delegation_help($route_name, RouteMatchInterface $route_match) {
   switch ($route_name) {
     case 'help.page':
-      $output = '';
       $output = '<p>' . t('This module allows site administrators to grant some roles the authority to assign selected roles to users, without them needing the <em>administer permissions</em> permission.') . '</p>';
       $output .= '<p>' . t('It provides its own tab in the user profile so that roles can be assigned without needing access to the user edit form.') . '</p>';
       return $output;
@@ -33,84 +30,19 @@ function role_delegation_help($route_name, RouteMatchInterface $route_match) {
 }
 
 /**
- * Implements hook_theme().
+ * Implements hook_ENTITY_TYPE_delete().
  */
-function role_delegation_theme() {
-  return array(
-    'role_delegation_delegate_roles_action_form' => array(
-      'render element' => 'form',
-    ),
-  );
-};
-
-/**
- * Implements hook_form_alter().
- */
-function role_delegation_form_alter(&$form, \Drupal\Core\Form\FormStateInterface &$form_state, $form_id) {
-  $current_user = \Drupal::currentUser();
-  if ($form_id == 'user_register_form' || $form_id == 'user_form') {
-    if ($current_user->hasPermission('administer permissions')) {
-      return;
-    }
-    $form = RoleDelegationHelper::addRoleDelegationElement($form, $form_state->getFormObject()->getEntity());
-    $form['#entity_builders'][] = 'role_delegation_update';
-  }
-  if ($form_id == 'views_form_user_admin_people_page_1') {
-    $options = $form['header']['user_bulk_form']['action']['#options'];
-    $roles = user_roles(TRUE);
-    unset($roles[AccountInterface::AUTHENTICATED_ROLE]);
-    unset($roles['administrator']);
-    foreach ($options as $option => $label) {
-      $operation_rid = explode('.', $option);
-      $operation = $operation_rid[0];
-      if ($operation == 'user_add_role_action') {
-        $rid = $operation_rid[1];
-        if (empty($roles[$rid]) || !$current_user->hasPermission("assign {$rid} role")){
-          unset($options[$option]);
-        }
-      }
-      elseif ($operation == 'user_remove_role_action') {
-        $rid = $operation_rid[1];
-        if (empty($roles[$rid]) || !$current_user->hasPermission("assign {$rid} role")){
-          unset($options[$option]);
-        }
-      }
-    }
-    $form['header']['user_bulk_form']['action']['#options'] = $options;
-  }
-}
-
-/**
- * Entity form builder to update roles to the user.
- */
-function role_delegation_update($entity_type, $entity, &$form, FormStateInterface $form_state) {
-  if($entity) {
-    foreach($form_state->getValue('roles_change') as $rid => $value) {
-      if($rid === $value) {
-        $entity->addRole($rid);
-      }
-      if($value === 0) {
-        $entity->removeRole($rid);
-      }
-    }
-  }
-}
-
-/**
- * Implements hook_entity_delete().
- */
-function role_delegation_entity_delete(Drupal\Core\Entity\EntityInterface $entity) {
-  if ($entity->getEntityTypeId() == 'user_role') {
-    $permission = "assign {$entity->id()} role";
-
-    $roles = \Drupal::entityQuery('user_role')
-      ->condition('permissions.*', $permission)
-      ->condition('id', $entity->id(), '<>')
-      ->execute();
-
-    foreach (Role::loadMultiple($roles) as $role) {
-      $role->revokePermission($permission);
-      $role->save();
-    }
+function role_delegation_user_role_delete(RoleInterface $entity) {
+  $permission = sprintf('assign %s role', $entity->id());
+
+  $roles = \Drupal::entityQuery('user_role')
+    ->condition('permissions.*', $permission)
+    ->condition('id', $entity->id(), '<>')
+    ->execute();
+
+  /** @var \Drupal\user\RoleInterface $role */
+  foreach (Role::loadMultiple($roles) as $role) {
+    $role->revokePermission($permission);
+    $role->save();
   }
 }
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..dc6a0a2 100644
--- a/role_delegation.services.yml
+++ b/role_delegation.services.yml
@@ -1,6 +1,8 @@
 services:
-  role_delegation.access_checker:
+  access_check.role_delegation:
     class: Drupal\role_delegation\Access\RoleDelegationAccessCheck
-    arguments: ['@current_user']
+    arguments: ['@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..a7cf422 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,23 @@ 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\role_delegation\RoleDelegationInterface $role_delegation
+   *   The role delegation service.
+   */
+  public function __construct(RoleDelegationInterface $role_delegation) {
+    $this->roleDelegation = $role_delegation;
+  }
+
+  /**
    * Custom access check for the /user/%/roles page.
    *
    * @param \Drupal\Core\Session\AccountInterface $account
@@ -28,8 +46,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..f551413 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']['role_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('role_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/src/Tests/RoleAssignTest.php b/src/Tests/RoleAssignTest.php
new file mode 100644
index 0000000..bf7834e
--- /dev/null
+++ b/src/Tests/RoleAssignTest.php
@@ -0,0 +1,110 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\role_delegation\Tests\RoleAssignTest.
+ */
+
+namespace Drupal\role_delegation\Tests;
+
+use Drupal\simpletest\WebTestBase;
+use Drupal\user\Entity\User;
+
+/**
+ * Functional tests for assigning roles.
+ *
+ * @group role_delegation
+ */
+class RoleAssignTest extends WebTestBase {
+
+  /**
+   * Modules to enable.
+   *
+   * @var array
+   */
+  public static $modules = ['user', 'role_delegation'];
+
+  /**
+   * The high role.
+   *
+   * @var string
+   */
+  protected $ridHigh;
+
+  /**
+   * The low role.
+   *
+   * @var string
+   */
+  protected $ridLow;
+
+  /**
+   * {@inheritdoc}
+   */
+  protected function setUp() {
+    parent::setUp();
+
+    // Create roles
+    $this->ridLow  = $this->drupalCreateRole([], 'low', 'low');
+    $this->ridHigh = $this->drupalCreateRole(['assign low role'], 'high', 'high');
+
+    // Create privileged user and log in
+    $this->drupalLogin($this->drupalCreateUser(['administer users', 'administer permissions']));
+  }
+
+  /**
+   * Ensure we can only see the roles we have permission to assign.
+   */
+  public function testRoleAccess() {
+    $rid1 = $this->drupalCreateRole([]);
+    $rid2 = $this->drupalCreateRole([]);
+    $rid3 = $this->drupalCreateRole([]);
+
+    $current_user = $this->drupalCreateUser(["assign $rid1 role", "assign $rid2 role"]);
+    $this->drupalLogin($current_user);
+
+    // Only 2 of the 3 roles appear on any user edit page.
+    $account = $this->drupalCreateUser();
+    $this->drupalGet(sprintf('/user/%s/roles', $account->id()));
+    $this->assertFieldByName("role_change[$rid1]");
+    $this->assertFieldByName("role_change[$rid2]");
+    $this->assertNoFieldByName("role_change[$rid3]");
+  }
+
+  /**
+   * Test that we can assign roles we have access to.
+   */
+  public function testRoleAssign() {
+    // Create a role and login as a user with the permission to assign it.
+    $rid1 = $this->drupalCreateRole([]);
+    $current_user = $this->drupalCreateUser(["assign $rid1 role"]);
+    $this->drupalLogin($current_user);
+
+    // Go to the users roles edit page.
+    $account = $this->drupalCreateUser();
+    $this->drupalGet(sprintf('/user/%s/roles', $account->id()));
+
+    // The form element field id and name.
+    $field_id = "edit-role-change-$rid1";
+    $field_name = "role_change[$rid1]";
+
+    // Ensure its disabled by default.
+    $this->assertNoFieldChecked($field_id, 'Role is not assigned by default.');
+    $this->assertFalse($account->hasPermission("assign $rid1 role"), 'Role is not assigned by default.');
+
+    // Assign the role and ensure its now checked and assigned.
+    $this->drupalPostForm(NULL, [$field_name => $rid1], 'Save');
+    $this->assertFieldChecked($field_id, 'Role has been granted.');
+    \Drupal::entityTypeManager()->clearCachedDefinitions();
+    $account = User::load($account->id());
+    $this->assertTrue($account->hasRole($rid1), 'Role has been granted');
+
+    // Revoke the role.
+    $this->drupalPostForm(NULL, [$field_name => FALSE], 'Save');
+    $this->assertNoFieldChecked($field_id, 'Role has been revoked.');
+    \Drupal::entityTypeManager()->clearCachedDefinitions();
+    $account = User::load($account->id());
+    $this->assertFalse($account->hasRole($rid1), 'Role has been revoked.');
+  }
+
+}
diff --git a/src/Tests/RoleDelegationEditingTest.php b/src/Tests/RoleDelegationEditingTest.php
deleted file mode 100644
index b1f0566..0000000
--- a/src/Tests/RoleDelegationEditingTest.php
+++ /dev/null
@@ -1,38 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\role_delegation\Tests\RoleDelegationEditingTest.
- */
-
-namespace Drupal\role_delegation\Tests;
-
-/**
- * Functional tests for editing roles.
- *
- * @group role_delegation
- */
-class RoleDelegationEditingTest extends RoleDelegationTestBase {
-
-  /**
-   * Rename a role, and check that users that had permission to assign
-   * the old role now have permission to assign the new one.
-   */
-  public function testRenameRole() {
-    $this->drupalPostForm("admin/people/roles/manage/{$this->ridLow}", array('label' => 'new low'), t('Save'));
-    $this->drupalGet('admin/people/permissions');
-    $this->assertFieldChecked("edit-{$this->ridHigh}-assign-low-role", 'Permissions are updated when role is renamed.');
-  }
-
-  /**
-   * Delete a role, then create a new one with the same name.
-   * Check that no users have permission to assign the new role.
-   */
-  public function testDeleteRole() {
-    $this->drupalPostForm("admin/people/roles/manage/{$this->ridLow}/delete", NULL, t('Delete'));
-    $this->drupalPostForm('admin/people/roles/add', array('label' => 'low', 'id' => 'low'), t('Save'));
-    $this->drupalGet('admin/people/permissions');
-    $this->assertNoFieldChecked("edit-{$this->ridHigh}-assign-low-role", 'Permissions are updated when role is deleted.');
-  }
-
-}
diff --git a/src/Tests/RoleDelegationOperationsTest.php b/src/Tests/RoleDelegationOperationsTest.php
deleted file mode 100644
index f2c52c0..0000000
--- a/src/Tests/RoleDelegationOperationsTest.php
+++ /dev/null
@@ -1,93 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\role_delegation\Tests\RoleDelegationOperationsTest.
- */
-
-namespace Drupal\role_delegation\Tests;
-
-/**
- * Functional tests for operations.
- *
- * @group role_delegation
- */
-class RoleDelegationOperationsTest extends RoleDelegationTestBase {
-
-  public static $modules = ['user', 'role_delegation', 'views'];
-
-  /**
-   * Check that the right combination of Add and Remove role
-   * operations is present in the user bulk update form.
-   */
-  public function testOperationsExist() {
-    $this->drupalLogin($this->userHigh);
-    $this->drupalGet('admin/people');
-
-    $this->assertFieldByXPath('//select[@name="action"]//option', "user_add_role_action.{$this->ridLow}", 'High user can use Add low role operation.');
-    $this->assertFieldByXPath('//select[@name="action"]//option', "user_remove_role_action.{$this->ridLow}", 'High user can use Remove low role operation.');
-    $this->assertNoFieldByXPath('//select[@name="action"]//option', "user_add_role_action.{$this->ridHigh}", 'High user can\'t use Add high role operation.');
-    $this->assertNoFieldByXPath('//select[@name="action"]//option', "user_remove_role_action.{$this->ridHigh}", 'High user can\'t use Remove high role operation.');
-  }
-
-  /**
-   * Check that Add and Remove role operations work as intended.
-   */
-  public function testOperationsWork() {
-    $uids_to_test = [$this->userHigh->uid, $this->userLow->uid];
-    $edit = [];
-    foreach ($uids_to_test as $uid) {
-      $userid = $uid->value;
-      $edit["user_bulk_form[$userid]"] = TRUE;
-    }
-
-    $this->drupalLogin($this->userHigh);
-    $this->drupalGet('admin/people');
-
-    // Add low role
-    $edit['action'] = "user_add_role_action.{$this->ridLow}";
-    $this->drupalPostForm(NULL, $edit, t('Apply'));
-    foreach ($uids_to_test as $uid) {
-      $userid = $uid->value;
-      $this->assertFieldByXPath("//tbody/tr[$userid+1]/td[4]//li", "{$this->ridLow}", sprintf('High user assigned low role to user %s', $userid));
-    }
-
-    // Remove low role
-    $edit['action'] = "user_remove_role_action.{$this->ridLow}";
-    $this->drupalPostForm(NULL, $edit, t('Apply'));
-    foreach ($uids_to_test as $uid) {
-      $userid = $uid->value;
-      $this->assertNoFieldByXPath("//tbody/tr[$userid]/td[4]//li", 'low', sprintf('High user removed low role from user %s.', $userid));
-    }
-  }
-
-  /**
-   * Check that operations can't be forged.
-   */
-  public function testOperationsForgery() {
-    $this->drupalLogin($this->userHigh);
-    $this->drupalGet('admin/people');
-
-    // Forge an operation to add the high role...
-    $option = $this->xpath("//select[@name='action']//option[@value='user_add_role_action.{$this->ridLow}']");
-    if (count($option) == 0) {
-      return;
-    }
-    $dome = dom_import_simplexml($option[0]);
-    $dome->setAttribute('value', "user_add_role_action.{$this->ridHigh}");
-
-    // Submit the form, and check that it wasn't granted.
-    $edit = [
-      "user_bulk_form[{$this->userLow->id()}]" => TRUE,
-      "action" => "user_add_role_action.{$this->ridHigh}",
-    ];
-    $this->drupalPostForm(NULL, $edit, t('Apply'));
-    $this->assertRaw(
-      t('An illegal choice has been detected. Please contact the site administrator.'),
-      t('Role assignment forgery is blocked.') . ' (#1)',
-      t('Role Delegation')
-    );
-    $this->assertNoFieldByXPath("//tbody/tr[{$this->userHigh->id()}]/td[4]//li", 'high', 'Role assignment forgery is blocked. (#2)');
-  }
-
-}
diff --git a/src/Tests/RoleDelegationPermissionsTest.php b/src/Tests/RoleDelegationPermissionsTest.php
deleted file mode 100644
index 9f683ed..0000000
--- a/src/Tests/RoleDelegationPermissionsTest.php
+++ /dev/null
@@ -1,53 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\role_delegation\Tests\RoleDelegationPermissionsTest.
- */
-
-namespace Drupal\role_delegation\Tests;
-
-/**
- * Functional tests for permissions.
- *
- * @group role_delegation
- */
-class RoleDelegationPermissionsTest extends RoleDelegationTestBase {
-
-  /**
-   * Check that high role can assign low role.
-   */
-  public function testHighLow() {
-    $this->drupalLogin($this->userHigh);
-    $this->assertTrue($this->assignRoleToUser($this->ridLow, $this->userHigh), sprintf("%s role can assign %s role.", 'High', 'low'));
-  }
-
-  /**
-   * Check that high role can't assign high role.
-   */
-  public function testHighHigh() {
-    $this->drupalLogin($this->userHigh);
-    // Just check that no option is presented to the user.
-    $this->assertFalse($this->assignRoleToUser($this->ridHigh, $this->userHigh), sprintf("%s role can't assign %s role.", 'High', 'high'));
-  }
-
-  /**
-   * Check that roles can't be assigned by forgery.
-   */
-  public function testRoleForgery() {
-    $this->drupalLogin($this->userHigh);
-
-    // Have the nefarious high user forge an option to assign the high role...
-    $this->drupalGet("user/{$this->userLow->id()}/edit");
-    $name = "roles_change[{$this->ridLow}]";
-    $input = $this->xpath("//input[@name='$name']");
-    $dom = dom_import_simplexml($input[0]);
-    $dom->setAttribute('value', $this->ridHigh);
-
-    // Submit the form, and check that he didn't get the role.
-    $this->drupalPostForm(NULL, array($name => TRUE), t('Save'));
-    $this->assertRaw('An illegal choice has been detected. Please contact the site administrator.', 'Role assignment forgery is blocked. (#1)');
-    $this->assertFieldByName($name, $this->ridLow, 'Role assignment forgery is blocked. (#2)', 'Role Delegation');
-  }
-
-}
diff --git a/src/Tests/RoleDelegationTestBase.php b/src/Tests/RoleDelegationTestBase.php
deleted file mode 100644
index 86bf02b..0000000
--- a/src/Tests/RoleDelegationTestBase.php
+++ /dev/null
@@ -1,135 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\role_delegation\Tests\RoleDelegationTestBase.
- */
-
-namespace Drupal\role_delegation\Tests;
-
-use Drupal\simpletest\WebTestBase;
-
-/**
- * Base class for Role Delegation tests.
- */
-abstract class RoleDelegationTestBase extends WebTestBase {
-
-  /**
-   * Modules to enable.
-   *
-   * @var array
-   */
-  public static $modules = ['user', 'role_delegation'];
-
-  /**
-   * The high role.
-   *
-   * @var string
-   */
-  protected $ridHigh;
-
-  /**
-   * The low role.
-   *
-   * @var string
-   */
-  protected $ridLow;
-
-  /**
-   * The high user.
-   *
-   * @var \Drupal\user\Entity\User
-   */
-  protected $userHigh;
-
-  /**
-   * The low user.
-   *
-   * @var \Drupal\user\Entity\User
-   */
-  protected $userLow;
-
-  /**
-   * {@inheritdoc}
-   */
-  protected function setUp() {
-    parent::setUp();
-
-    // Create roles
-    $this->ridHigh = $this->drupalCreateRole(array(), 'high', 'high');
-    $this->ridLow  = $this->drupalCreateRole(array(), 'low', 'low');
-
-    // Create users
-    $this->userHigh = $this->drupalCreateUser(array('administer users'));
-    $this->userLow  = $this->drupalCreateUser(array('administer users'));
-
-    // Create privileged user and log in
-    $this->drupalLogin($this->drupalCreateUser(array('administer users', 'administer permissions')));
-
-    // Assign permissions to roles
-    $this->assignPermissionToRole('assign low role', $this->ridHigh);  // 'high' can assign 'low'
-
-    // Assign roles to users
-    $this->assertTrue($this->assignRoleToUser($this->ridHigh, $this->userHigh), 'Assign high role to high user');
-  }
-
-  /**
-   * Assign or remove one role to/from one user.
-   *
-   * The logged in user must have the "administer users"
-   * permission in order for this function to succeed.
-   *
-   * @param $rid
-   *   The role id of the role to assign or remove.
-   * @param $user
-   *   The user object of the user to assign/remove the role to.
-   * @param $assign
-   *   TRUE (the default) to assign the role, or
-   *   FALSE to remove it.
-   *
-   * @return bool
-   *   TRUE or FALSE depending on whether the permission was successfully
-   *   assigned or removed.
-   */
-  protected function assignRoleToUser($rid, $user, $assign = TRUE) {
-    $this->drupalGet("user/{$user->id()}/edit");
-
-    if (count($this->xpath("//input[@name='roles[$rid]']"))) {
-      $name = "roles[$rid]";
-    }
-    elseif (count($this->xpath("//input[@name='roles_change[$rid]']"))) {
-      $name = "roles_change[$rid]";
-    }
-    else {
-      return FALSE;
-    }
-    $this->drupalPostForm(NULL, array($name => $assign), t('Save'));
-
-    $elements = $this->xpath("//input[@name='$name']");
-    return isset($elements[0]) && ($assign XOR empty($elements[0]['checked']));
-  }
-
-  /**
-   * Assign or remove one permission to/from one role, and assert
-   * that the result succeeded.
-   *
-   * @param $permission
-   *   The name of the permission to assign or remove.
-   * @param $rid
-   *   The role id of the role to assign/remove the permission to/from.
-   * @param $assign
-   *   TRUE (the default) to assign the permission, or
-   *   FALSE to remove it.
-   *
-   * @return bool
-   *   TRUE or FALSE depending on whether the permission was
-   *   successfully assigned or removed.
-   */
-  protected function assignPermissionToRole($permission, $rid, $assign = TRUE) {
-    $name = "{$rid}[{$permission}]";
-    $this->drupalPostForm("admin/people/permissions/$rid", array($name => $assign), t('Save permissions'));
-    $elements = $this->xpath("//input[@name='$name']");
-    $this->assertTrue(isset($elements[0]) && ($assign XOR empty($elements[0]['checked'])), ($assign ? 'Assign' : 'Remove') . ' permission "' . $permission . '" ' . ($assign ? 'to' : 'from') . " role $rid.");
-  }
-
-}
diff --git a/tests/Kernel/AccessTest.php b/tests/Kernel/AccessTest.php
index 68f5488..cd3cff9 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;
 
@@ -38,7 +38,7 @@ class AccessTest extends KernelTestBase {
     parent::setUp();
     $this->installSchema('system', 'sequences');
     $this->installEntitySchema('user');
-    $this->accessChecker = $this->container->get('role_delegation.access_checker');
+    $this->accessChecker = $this->container->get('access_check.role_delegation');
 
     // User 1 is still a super user so we create that user first so moving
     // forward we're just using normal users.
@@ -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.
diff --git a/tests/Kernel/RoleDelegationTest.php b/tests/Kernel/RoleDelegationTest.php
new file mode 100644
index 0000000..4c3000d
--- /dev/null
+++ b/tests/Kernel/RoleDelegationTest.php
@@ -0,0 +1,94 @@
+<?php
+
+namespace Drupal\Tests\role_delegation\Kernel;
+
+use Drupal\KernelTests\KernelTestBase;
+use Drupal\simpletest\UserCreationTrait;
+use Drupal\user\Entity\Role;
+
+/**
+ * @coversDefaultClass \Drupal\role_delegation\RoleDelegation
+ */
+class RoleDelegationTest extends KernelTestBase {
+
+  use UserCreationTrait;
+
+  /**
+   * The modules to enable for this test.
+   *
+   * @var array
+   */
+  public static $modules = ['system', 'role_delegation', 'user'];
+
+  /**
+   * The Role Delegation service.
+   *
+   * @var \Drupal\role_delegation\RoleDelegationInterface
+   */
+  protected $roleDelegation;
+
+  /**
+   * {@inheritdoc}
+   */
+  protected function setUp() {
+    parent::setUp();
+    $this->installSchema('system', 'sequences');
+    $this->installEntitySchema('user');
+    $this->roleDelegation = $this->container->get('role_delegation');
+
+    // User 1 is still a super user so we create that user first so moving
+    // forward we're just using normal users.
+    $this->createUser([]);
+  }
+
+  /**
+   * Test the roles that can be assigned by a given user.
+   *
+   * @covers ::getAssignableRoles
+   */
+  public function testAssignableRoles() {
+    $rid1 = $this->createRole([]);
+    $rid2 = $this->createRole([]);
+    $rid3 = $this->createRole([]);
+
+    // Test the 'assign all roles permission'. We have to merge in the roles of
+    // the account as well because createUser() creates a new role.
+    $account = $this->createUser(['assign all roles']);
+    $this->assertEquals(array_merge([$rid1, $rid2, $rid3], $account->getRoles(TRUE)), array_keys($this->roleDelegation->getAssignableRoles($account)));
+
+    // If they have these two roles, they can assign exactly those two roles.
+    $account = $this->createUser(["assign $rid1 role", "assign $rid2 role"]);
+    $this->assertEquals([$rid1, $rid2], array_keys($this->roleDelegation->getAssignableRoles($account)));
+
+    // Doesn't matter what permissions they have here, they can never assign
+    // anonymous or authenticated roles.
+    $account = $this->createUser(['administer users', 'administer permissions']);
+    $this->assertEquals([], $this->roleDelegation->getAssignableRoles($account));
+  }
+
+  /**
+   * Test the all roles methods filters special roles.
+   *
+   * @covers ::getAllRoles
+   */
+  public function testGetAllRoles() {
+    $rid1 = $this->createRole([]);
+    $rid2 = $this->createRole([]);
+    $this->assertEquals([$rid1, $rid2], array_keys($this->roleDelegation->getAllRoles()));
+  }
+
+  /**
+   * Deleting a role revokes the permission allowing users to assign the role.
+   */
+  public function testDeleteRole() {
+    $rid = $this->createRole([]);
+    $permission = "assign $rid role";
+    $account = $this->createUser([$permission]);
+    $this->assertTrue($account->hasPermission($permission), sprintf('User has "%s" permission.', $permission));
+
+    // Delete the role and ensure the user no longer has the permission.
+    Role::load($rid)->delete();
+    $this->assertFalse($account->hasPermission($permission), sprintf('User no longer has "%s" permission.', $permission));
+  }
+
+}
