Removed generic getStorage method and replaced with distinct calls to getters

src/AbstractServer.php

@@ -13,6 +13,12 @@ namespace League\OAuth2\Server;
 
 use League\OAuth2\Server\Exception;
 use League\OAuth2\Server\TokenType\TokenTypeInterface;
+use League\OAuth2\Server\Storage\SessionInterface;
+use League\OAuth2\Server\Storage\AccessTokenInterface;
+use League\OAuth2\Server\Storage\RefreshTokenInterface;
+use League\OAuth2\Server\Storage\AuthCodeInterface;
+use League\OAuth2\Server\Storage\ScopeInterface;
+use League\OAuth2\Server\Storage\ClientInterface;
 use Symfony\Component\HttpFoundation\Request;
 use League\Event\Emitter;
 
@@ -30,10 +36,40 @@ abstract class AbstractServer
     protected $request;
 
     /**
-     * Storage classes
+     * Session storage
-     * @var array
+     * @var \League\OAuth2\Server\Storage\SessionInterface
      */
-    protected $storages = [];
+    protected $sessionStorage;
+
+    /**
+     * Access token storage
+     * @var \League\OAuth2\Server\Storage\AccessTokenInterface
+     */
+    protected $accessTokenStorage;
+
+    /**
+     * Refresh token storage
+     * @var \League\OAuth2\Server\Storage\RefreshTokenInterface
+     */
+    protected $refreshTokenStorage;
+
+    /**
+     * Auth code storage
+     * @var \League\OAuth2\Server\Storage\AuthCodeInterface
+     */
+    protected $authCodeStorage;
+
+    /**
+     * Scope storage
+     * @var \League\OAuth2\Server\Storage\ScopeInterface
+     */
+    protected $scopeStorage;
+
+    /**
+     * Client storage
+     * @var \League\OAuth2\Server\Storage\ClientInterface
+     */
+    protected $clientStorage;
 
     /**
      * Token type
@@ -113,19 +149,135 @@ abstract class AbstractServer
     }
 
     /**
-     * Return a storage class
+     * Set the client storage
-     * @param  string                                                                  $obj The class required
+     * @param  \League\OAuth2\Server\Storage\ClientInterface $storage
-     * @return Storage\ClientInterface|Storage\ScopeInterface|Storage\SessionInterface|Storage\AccessTokenInterface|Storage\AuthCodeInterface|Storage\RefreshTokenInterface
+     * @return self
      */
-    public function getStorage($obj)
+    public function setClientStorage(ClientInterface $storage)
     {
-        if (!isset($this->storages[$obj])) {
+        $storage->setServer($this);
-            throw new Exception\ServerErrorException(
+        $this->clientStorage = $storage;
-                'The `'.$obj.'` storage interface has not been registered with the server'
+
-            );
+        return $this;
     }
 
-        return $this->storages[$obj];
+    /**
+     * Set the session storage
+     * @param  \League\OAuth2\Server\Storage\SessionInterface $storage
+     * @return self
+     */
+    public function setSessionStorage(SessionInterface $storage)
+    {
+        $storage->setServer($this);
+        $this->sessionStorage = $storage;
+
+        return $this;
+    }
+
+    /**
+     * Set the access token storage
+     * @param  \League\OAuth2\Server\Storage\AccessTokenInterface $storage
+     * @return self
+     */
+    public function setAccessTokenStorage(AccessTokenInterface $storage)
+    {
+        $storage->setServer($this);
+        $this->accessTokenStorage = $storage;
+
+        return $this;
+    }
+
+    /**
+     * Set the refresh token storage
+     * @param  \League\OAuth2\Server\Storage\RefreshTokenInteface $storage
+     * @return self
+     */
+    public function setRefreshTokenStorage(RefreshTokenInterface $storage)
+    {
+        $storage->setServer($this);
+        $this->refreshTokenStorage = $storage;
+
+        return $this;
+    }
+
+    /**
+     * Set the auth code storage
+     * @param  \League\OAuth2\Server\Storage\AuthCodeInterface $authCode
+     * @return self
+     */
+    public function setAuthCodeStorage(AuthCodeInterface $storage)
+    {
+        $storage->setServer($this);
+        $this->authCodeStorage = $storage;
+
+        return $this;
+    }
+
+    /**
+     * Set the scope storage
+     * @param  \League\OAuth2\Server\Storage\ScopeInterface $storage
+     * @return self
+     */
+    public function setScopeStorage(ScopeInterface $storage)
+    {
+        $storage->setServer($this);
+        $this->scopeStorage = $storage;
+
+        return $this;
+    }
+
+    /**
+     * Return the client storage
+     * @return \League\OAuth2\Server\Storage\ClientInterface
+     */
+    public function getClientStorage()
+    {
+        return $this->clientStorage;
+    }
+
+    /**
+     * Return the scope storage
+     * @return \League\OAuth2\Server\Storage\ScopeInterface
+     */
+    public function getScopeStorage()
+    {
+        return $this->scopeStorage;
+    }
+
+    /**
+     * Return the session storage
+     * @return \League\OAuth2\Server\Storage\SessionInterface
+     */
+    public function getSessionStorage()
+    {
+        return $this->sessionStorage;
+    }
+
+    /**
+     * Return the refresh token storage
+     * @return \League\OAuth2\Server\Storage\RefreshTokenInterface
+     */
+    public function getRefreshTokenStorage()
+    {
+        return $this->refreshTokenStorage;
+    }
+
+    /**
+     * Return the access token storage
+     * @return \League\OAuth2\Server\Storage\AccessTokenInterface
+     */
+    public function getAccessTokenStorage()
+    {
+        return $this->accessTokenStorage;
+    }
+
+    /**
+     * Return the auth code storage
+     * @return \League\OAuth2\Server\Storage\AuthCodeInterface
+     */
+    public function getAuthCodeStorage()
+    {
+        return $this->authCodeStorage;
     }
 
     /**

src/AuthorizationServer.php

@@ -75,8 +75,6 @@ class AuthorizationServer extends AbstractServer
      */
     public function __construct()
     {
-        $this->storages = [];
-
         // Set Bearer as the default token type
         $this->setTokenType(new Bearer);
 
@@ -85,84 +83,6 @@ class AuthorizationServer extends AbstractServer
         return $this;
     }
 
-    /**
-     * Set the client storage
-     * @param  ClientInterface $storage
-     * @return self
-     */
-    public function setClientStorage(ClientInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['client'] = $storage;
-
-        return $this;
-    }
-
-    /**
-     * Set the session storage
-     * @param  SessionInterface $storage
-     * @return self
-     */
-    public function setSessionStorage(SessionInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['session'] = $storage;
-
-        return $this;
-    }
-
-    /**
-     * Set the access token storage
-     * @param  AccessTokenInterface $storage
-     * @return self
-     */
-    public function setAccessTokenStorage(AccessTokenInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['access_token'] = $storage;
-
-        return $this;
-    }
-
-    /**
-     * Set the refresh token storage
-     * @param  RefreshTokenInteface $storage
-     * @return self
-     */
-    public function setRefreshTokenStorage(RefreshTokenInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['refresh_token'] = $storage;
-
-        return $this;
-    }
-
-    /**
-     * Set the auth code storage
-     * @param  AuthCodeInterface $authCode
-     * @return self
-     */
-    public function setAuthCodeStorage(AuthCodeInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['auth_code'] = $storage;
-
-        return $this;
-    }
-
-    /**
-     * Set the scope storage
-     * @param  ScopeInterface $storage
-     * @return self
-     */
-    public function setScopeStorage(ScopeInterface $storage)
-    {
-        $storage->setServer($this);
-        $this->storages['scope'] = $storage;
-
-        return $this;
-    }
-
     /**
      * Enable support for a grant
      * @param  GrantTypeInterface $grantType  A grant class which conforms to Interface/GrantTypeInterface

src/Entity/AccessTokenEntity.php

@@ -26,7 +26,7 @@ class AccessTokenEntity extends AbstractTokenEntity
             return $this->session;
         }
 
-        $this->session = $this->server->getStorage('session')->getByAccessToken($this);
+        $this->session = $this->server->getSessionStorage()->getByAccessToken($this);
 
         return $this->session;
     }
@@ -53,7 +53,7 @@ class AccessTokenEntity extends AbstractTokenEntity
     {
         if ($this->scopes === null) {
             $this->scopes = $this->formatScopes(
-                $this->server->getStorage('access_token')->getScopes($this)
+                $this->server->getAccessTokenStorage()->getScopes($this)
             );
         }
 
@@ -65,7 +65,7 @@ class AccessTokenEntity extends AbstractTokenEntity
      */
     public function save()
     {
-        $this->server->getStorage('access_token')->create(
+        $this->server->getAccessTokenStorage()->create(
             $this->getId(),
             $this->getExpireTime(),
             $this->getSession()->getId()
@@ -73,7 +73,7 @@ class AccessTokenEntity extends AbstractTokenEntity
 
         // Associate the scope with the token
         foreach ($this->getScopes() as $scope) {
-            $this->server->getStorage('access_token')->associateScope($this, $scope);
+            $this->server->getAccessTokenStorage()->associateScope($this, $scope);
         }
 
         return $this;
@@ -84,6 +84,6 @@ class AccessTokenEntity extends AbstractTokenEntity
      */
     public function expire()
     {
-        $this->server->getStorage('access_token')->delete($this);
+        $this->server->getAccessTokenStorage()->delete($this);
     }
 }

src/Entity/AuthCodeEntity.php

@@ -70,7 +70,7 @@ class AuthCodeEntity extends AbstractTokenEntity
             return $this->session;
         }
 
-        $this->session = $this->server->getStorage('session')->getByAuthCode($this);
+        $this->session = $this->server->getSessionStorage()->getByAuthCode($this);
 
         return $this->session;
     }
@@ -83,7 +83,7 @@ class AuthCodeEntity extends AbstractTokenEntity
     {
         if ($this->scopes === null) {
             $this->scopes = $this->formatScopes(
-                $this->server->getStorage('auth_code')->getScopes($this)
+                $this->server->getAuthCodeStorage()->getScopes($this)
             );
         }
 
@@ -95,7 +95,7 @@ class AuthCodeEntity extends AbstractTokenEntity
      */
     public function save()
     {
-        $this->server->getStorage('auth_code')->create(
+        $this->server->getAuthCodeStorage()->create(
             $this->getId(),
             $this->getExpireTime(),
             $this->getSession()->getId(),
@@ -104,7 +104,7 @@ class AuthCodeEntity extends AbstractTokenEntity
 
         // Associate the scope with the token
         foreach ($this->getScopes() as $scope) {
-            $this->server->getStorage('auth_code')->associateScope($this, $scope);
+            $this->server->getAuthCodeStorage()->associateScope($this, $scope);
         }
 
         return $this;
@@ -115,6 +115,6 @@ class AuthCodeEntity extends AbstractTokenEntity
      */
     public function expire()
     {
-        $this->server->getStorage('auth_code')->delete($this);
+        $this->server->getAuthCodeStorage()->delete($this);
     }
 }

src/Entity/RefreshTokenEntity.php

@@ -59,7 +59,7 @@ class RefreshTokenEntity extends AbstractTokenEntity
     public function getAccessToken()
     {
         if (! $this->accessTokenEntity instanceof AccessTokenEntity) {
-            $this->accessTokenEntity = $this->server->getStorage('access_token')->get($this->accessTokenId);
+            $this->accessTokenEntity = $this->server->getAccessTokenStorage()->get($this->accessTokenId);
         }
 
         return $this->accessTokenEntity;
@@ -70,7 +70,7 @@ class RefreshTokenEntity extends AbstractTokenEntity
      */
     public function save()
     {
-        $this->server->getStorage('refresh_token')->create(
+        $this->server->getRefreshTokenStorage()->create(
             $this->getId(),
             $this->getExpireTime(),
             $this->getAccessToken()->getId()
@@ -82,6 +82,6 @@ class RefreshTokenEntity extends AbstractTokenEntity
      */
     public function expire()
     {
-        $this->server->getStorage('refresh_token')->delete($this);
+        $this->server->getRefreshTokenStorage()->delete($this);
     }
 }

src/Entity/SessionEntity.php

@@ -142,7 +142,7 @@ class SessionEntity
     public function getScopes()
     {
         if ($this->scopes === null) {
-            $this->scopes = $this->formatScopes($this->server->getStorage('session')->getScopes($this));
+            $this->scopes = $this->formatScopes($this->server->getSessionStorage()->getScopes($this));
         }
 
         return $this->scopes;
@@ -213,7 +213,7 @@ class SessionEntity
             return $this->client;
         }
 
-        $this->client = $this->server->getStorage('client')->getBySession($this);
+        $this->client = $this->server->getClientStorage()->getBySession($this);
 
         return $this->client;
     }
@@ -259,7 +259,7 @@ class SessionEntity
     public function save()
     {
         // Save the session and get an identifier
-        $id = $this->server->getStorage('session')->create(
+        $id = $this->server->getSessionStorage()->create(
             $this->getOwnerType(),
             $this->getOwnerId(),
             $this->getClient()->getId(),
@@ -270,7 +270,7 @@ class SessionEntity
 
         // Associate the scope with the session
         foreach ($this->getScopes() as $scope) {
-            $this->server->getStorage('session')->associateScope($this, $scope);
+            $this->server->getSessionStorage()->associateScope($this, $scope);
         }
     }
 }

src/Grant/AbstractGrant.php

@@ -153,7 +153,7 @@ abstract class AbstractGrant implements GrantTypeInterface
         $scopes = [];
 
         foreach ($scopesList as $scopeItem) {
-            $scope = $this->server->getStorage('scope')->get(
+            $scope = $this->server->getScopeStorage()->get(
                 $scopeItem,
                 $this->getIdentifier(),
                 $client->getId()

src/Grant/AuthCodeGrant.php

@@ -85,7 +85,7 @@ class AuthCodeGrant extends AbstractGrant
         }
 
         // Validate client ID and redirect URI
-        $client = $this->server->getStorage('client')->get(
+        $client = $this->server->getClientStorage()->get(
             $clientId,
             null,
             $redirectUri,
@@ -186,7 +186,7 @@ class AuthCodeGrant extends AbstractGrant
         }
 
         // Validate client ID and client secret
-        $client = $this->server->getStorage('client')->get(
+        $client = $this->server->getClientStorage()->get(
             $clientId,
             $clientSecret,
             $redirectUri,
@@ -204,7 +204,7 @@ class AuthCodeGrant extends AbstractGrant
             throw new Exception\InvalidRequestException('code');
         }
 
-        $code = $this->server->getStorage('auth_code')->get($authCode);
+        $code = $this->server->getAuthCodeStorage()->get($authCode);
         if (($code instanceof AuthCodeEntity) === false) {
             throw new Exception\InvalidRequestException('code');
         }

src/Grant/ClientCredentialsGrant.php

@@ -71,7 +71,7 @@ class ClientCredentialsGrant extends AbstractGrant
         }
 
         // Validate client ID and client secret
-        $client = $this->server->getStorage('client')->get(
+        $client = $this->server->getClientStorage()->get(
             $clientId,
             $clientSecret,
             null,

src/Grant/PasswordGrant.php

@@ -95,7 +95,7 @@ class PasswordGrant extends AbstractGrant
         }
 
         // Validate client ID and client secret
-        $client = $this->server->getStorage('client')->get(
+        $client = $this->server->getClientStorage()->get(
             $clientId,
             $clientSecret,
             null,

src/Grant/RefreshTokenGrant.php

@@ -76,7 +76,7 @@ class RefreshTokenGrant extends AbstractGrant
         }
 
         // Validate client ID and client secret
-        $client = $this->server->getStorage('client')->get(
+        $client = $this->server->getClientStorage()->get(
             $clientId,
             $clientSecret,
             null,
@@ -94,7 +94,7 @@ class RefreshTokenGrant extends AbstractGrant
         }
 
         // Validate refresh token
-        $oldRefreshToken = $this->server->getStorage('refresh_token')->get($oldRefreshTokenParam);
+        $oldRefreshToken = $this->server->getRefreshTokenStorage()->get($oldRefreshTokenParam);
 
         if (($oldRefreshToken instanceof RefreshTokenEntity) === false) {
             throw new Exception\InvalidRefreshException();
@@ -136,7 +136,7 @@ class RefreshTokenGrant extends AbstractGrant
         }
 
         // Expire the old token and save the new one
-        $oldAccessToken->expire($this->server->getStorage('access_token'));
+        $oldAccessToken->expire();
         $newAccessToken->save();
 
         $this->server->getTokenType()->setSession($session);

src/ResourceServer.php

@@ -27,7 +27,7 @@ class ResourceServer extends AbstractServer
 {
     /**
      * The access token
-     * @var League\OAuth2\Server\AccessToken
+     * @var \League\OAuth2\Server\Entity\AccessTokenEntity
      */
     protected $accessToken;
 
@@ -51,17 +51,10 @@ class ResourceServer extends AbstractServer
         ClientInterface $clientStorage,
         ScopeInterface $scopeStorage
     ) {
-        $sessionStorage->setServer($this);
+        $this->setSessionStorage($sessionStorage);
-        $this->setStorage('session', $sessionStorage);
+        $this->setAccessTokenStorage($accessTokenStorage);
-
+        $this->setClientStorage($clientStorage);
-        $accessTokenStorage->setServer($this);
+        $this->setScopeStorage($scopeStorage);
-        $this->setStorage('access_token', $accessTokenStorage);
-
-        $clientStorage->setServer($this);
-        $this->setStorage('client', $clientStorage);
-
-        $scopeStorage->setServer($this);
-        $this->setStorage('scope', $scopeStorage);
 
         // Set Bearer as the default token type
         $this->setTokenType(new Bearer);
@@ -71,20 +64,6 @@ class ResourceServer extends AbstractServer
         return $this;
     }
 
-    /**
-     * Set the storage
-     * @param  string $type    Storage type
-     * @param  mixed  $storage Storage class
-     * @return self
-     */
-    protected function setStorage($type, $storage)
-    {
-        $storage->setServer($this);
-        $this->storages[$type] = $storage;
-
-        return $this;
-    }
-
     /**
      * Returns the query string key for the access token.
      * @return string
@@ -185,7 +164,7 @@ class ResourceServer extends AbstractServer
                                 : $this->determineAccessToken($headersOnly);
 
         // Set the access token
-        $this->accessToken = $this->storages['access_token']->get($accessTokenString);
+        $this->accessToken = $this->getAccessTokenStorage()->get($accessTokenString);
 
         if (!$this->accessToken instanceof AccessTokenEntity) {
             throw new Exception\AccessDeniedException;

tests/unit/AbstractServerTest.php

@@ -25,11 +25,4 @@ class AbstractServerTest extends \PHPUnit_Framework_TestCase
         $this->assertTrue($server2->getRequest() instanceof \Symfony\Component\HttpFoundation\Request);
 
     }
-
-    public function testGetStorageException()
-    {
-        $this->setExpectedException('League\OAuth2\Server\Exception\ServerErrorException');
-        $server = new StubAbstractServer();
-        $server->getStorage('foobar');
-    }
 }

tests/unit/AuthorizationServerTest.php

@@ -34,7 +34,7 @@ class AuthorizationServerTest extends \PHPUnit_Framework_TestCase
         $this->assertSame($server->getResponseTypes(), ['foobar']);
         $this->assertTrue($server->scopeParamRequired());
         $this->assertTrue($server->stateParamRequired());
-        $this->assertTrue($server->getStorage('scope') instanceof ScopeInterface);
+        $this->assertTrue($server->getScopeStorage() instanceof ScopeInterface);
         $this->assertEquals('foobar', $server->getDefaultScope());
         $this->assertEquals(',', $server->getScopeDelimeter());
         $this->assertEquals(1, $server->getAccessTokenTTL());

tests/unit/Entity/AccessTokenEntityTest.php

@@ -29,8 +29,8 @@ class AccessTokenTest extends \PHPUnit_Framework_TestCase
         );
         $sessionStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
-        $server->shouldReceive('getStorage')->with('access_token')->andReturn($accessTokenStorage);
+        $server->shouldReceive('getAccessTokenStorage')->andReturn($accessTokenStorage);
 
         $server->setAccessTokenStorage($accessTokenStorage);
         $server->setSessionStorage($sessionStorage);
@@ -49,7 +49,7 @@ class AccessTokenTest extends \PHPUnit_Framework_TestCase
         $accessTokenStorage->shouldReceive('delete');
         $accessTokenStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('access_token')->andReturn($accessTokenStorage);
+        $server->shouldReceive('getAccessTokenStorage')->andReturn($accessTokenStorage);
 
         $server->setAccessTokenStorage($accessTokenStorage);
 

tests/unit/Entity/AuthCodeEntityTest.php

@@ -40,7 +40,7 @@ class AuthCodeTest extends \PHPUnit_Framework_TestCase
             (new ScopeEntity($server))->hydrate(['id' => 'foo'])
         ]);
 
-        $server->shouldReceive('getStorage')->with('auth_code')->andReturn($authCodeStorage);
+        $server->shouldReceive('getAuthCodeStorage')->andReturn($authCodeStorage);
 
         $sessionStorage = M::mock('League\OAuth2\Server\Storage\SessionInterface');
         $sessionStorage->shouldReceive('getByAuthCode')->andReturn(
@@ -48,7 +48,7 @@ class AuthCodeTest extends \PHPUnit_Framework_TestCase
         );
         $sessionStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
 
         $server->setAuthCodeStorage($authCodeStorage);
         $server->setSessionStorage($sessionStorage);

tests/unit/Entity/RefreshTokenEntityTest.php

@@ -47,7 +47,7 @@ class RefreshTokenEntityTest extends \PHPUnit_Framework_TestCase
         $refreshTokenStorage->shouldReceive('setServer');
         $refreshTokenStorage->shouldReceive('associateScope');
 
-        $server->shouldReceive('getStorage')->with('refresh_token')->andReturn($refreshTokenStorage);
+        $server->shouldReceive('getRefreshTokenStorage')->andReturn($refreshTokenStorage);
 
         $accessTokenStorage = M::mock('League\OAuth2\Server\Storage\AccessTokenInterface');
         $accessTokenStorage->shouldReceive('setServer');
@@ -58,7 +58,7 @@ class RefreshTokenEntityTest extends \PHPUnit_Framework_TestCase
             (new ScopeEntity($server))->hydrate(['id' => 'foo'])
         ]);
 
-        $server->shouldReceive('getStorage')->with('access_token')->andReturn($accessTokenStorage);
+        $server->shouldReceive('getAccessTokenStorage')->andReturn($accessTokenStorage);
 
         $sessionStorage = M::mock('League\OAuth2\Server\Storage\SessionInterface');
         $sessionStorage->shouldReceive('getByAccessToken')->andReturn(
@@ -66,7 +66,7 @@ class RefreshTokenEntityTest extends \PHPUnit_Framework_TestCase
         );
         $sessionStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
 
         $server->setAccessTokenStorage($accessTokenStorage);
         $server->setRefreshTokenStorage($refreshTokenStorage);
@@ -84,7 +84,7 @@ class RefreshTokenEntityTest extends \PHPUnit_Framework_TestCase
         $refreshTokenStorage->shouldReceive('delete');
         $refreshTokenStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('refresh_token')->andReturn($refreshTokenStorage);
+        $server->shouldReceive('getRefreshTokenStorage')->andReturn($refreshTokenStorage);
 
         $server->setRefreshTokenStorage($refreshTokenStorage);
 

tests/unit/Entity/SessionEntityTest.php

@@ -81,7 +81,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
         $accessTokenStorage->shouldReceive('setServer');
         $server->setAccessTokenStorage($accessTokenStorage);
 
-        $server->shouldReceive('getStorage')->with('access_token')->andReturn($accessTokenStorage);
+        $server->shouldReceive('getAccessTokenStorage')->andReturn($accessTokenStorage);
 
         $sessionStorage = M::mock('League\OAuth2\Server\Storage\SessionInterface');
         $sessionStorage->shouldReceive('getScopes')->andReturn(
@@ -90,7 +90,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
         $sessionStorage->shouldReceive('setServer');
         $server->setSessionStorage($sessionStorage);
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
 
         $entity = new SessionEntity($server);
         $this->assertEquals($entity->getScopes(), []);
@@ -106,7 +106,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
         $accessTokenStorage->shouldReceive('setServer');
         $server->setAccessTokenStorage($accessTokenStorage);
 
-        $server->shouldReceive('getStorage')->with('access_token')->andReturn($accessTokenStorage);
+        $server->shouldReceive('getAccessTokenStorage')->andReturn($accessTokenStorage);
 
         $sessionStorage = M::mock('League\OAuth2\Server\Storage\SessionInterface');
         $sessionStorage->shouldReceive('getScopes')->andReturn(
@@ -115,7 +115,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
         $sessionStorage->shouldReceive('setServer');
         $server->setSessionStorage($sessionStorage);
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
 
         $entity = new SessionEntity($server);
         $this->assertFalse($entity->hasScope('foo'));
@@ -135,7 +135,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
             (new ScopeEntity($server))->hydrate(['id' => 'foo'])
         ]);
 
-        $server->shouldReceive('getStorage')->with('session')->andReturn($sessionStorage);
+        $server->shouldReceive('getSessionStorage')->andReturn($sessionStorage);
 
         $clientStorage = M::mock('League\OAuth2\Server\Storage\ClientInterface');
         $clientStorage->shouldReceive('getBySession')->andReturn(
@@ -143,7 +143,7 @@ class SessionTest extends \PHPUnit_Framework_TestCase
         );
         $clientStorage->shouldReceive('setServer');
 
-        $server->shouldReceive('getStorage')->with('client')->andReturn($clientStorage);
+        $server->shouldReceive('getClientStorage')->andReturn($clientStorage);
 
         $server->setSessionStorage($sessionStorage);
         $server->setClientStorage($clientStorage);