diff --git a/src/Tests/TokenUnitTest.php b/src/Tests/TokenUnitTest.php
index 0309a43..30e766d 100644
--- a/src/Tests/TokenUnitTest.php
+++ b/src/Tests/TokenUnitTest.php
@@ -14,6 +14,11 @@ namespace Drupal\token\Tests;
 class TokenUnitTest extends TokenKernelTestBase {
 
   /**
+   * @var \Drupal\token\Token
+   */
+  protected $tokenService;
+
+  /**
    * Modules to enable.
    *
    * @var array
@@ -21,6 +26,15 @@ class TokenUnitTest extends TokenKernelTestBase {
   public static $modules = array('file', 'node', 'token_test');
 
   /**
+   * {@inheritdoc}
+   */
+  public function setUp() {
+    parent::setUp();
+
+    $this->tokenService = \Drupal::token();
+  }
+
+  /**
    * Test token_get_invalid_tokens() and token_get_invalid_tokens_by_context().
    */
   public function testGetInvalidTokens() {
@@ -90,7 +104,7 @@ class TokenUnitTest extends TokenKernelTestBase {
       $tokens = array_merge($test['valid tokens'], $test['invalid tokens']);
       shuffle($tokens);
 
-      $invalid_tokens = token_get_invalid_tokens_by_context(implode(' ', $tokens), $test['types']);
+      $invalid_tokens = $this->tokenService->getInvalidTokensByContext(implode(' ', $tokens), $test['types']);
 
       sort($invalid_tokens);
       sort($test['invalid tokens']);