From 07a42f6f439147e39e0ea78dc014188d1f860d6b Mon Sep 17 00:00:00 2001
From: Alex Bilbie <hello@alexbilbie.com>
Date: Wed, 6 Aug 2014 08:42:42 +0100
Subject: [PATCH] Added setAccessTokenId method

---
 src/Entity/RefreshTokenEntity.php            | 30 ++++++++++++++++----
 tests/unit/Entity/RefreshTokenEntityTest.php | 19 +++++++++++--
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/Entity/RefreshTokenEntity.php b/src/Entity/RefreshTokenEntity.php
index 7ed35a76..070fddf4 100644
--- a/src/Entity/RefreshTokenEntity.php
+++ b/src/Entity/RefreshTokenEntity.php
@@ -20,16 +20,34 @@ class RefreshTokenEntity extends AbstractTokenEntity
      * Access token associated to refresh token
      * @var \League\OAuth2\Server\Entity\AccessTokenEntity
      */
-    protected $accessToken;
+    protected $accessTokenEntity;
+
+    /**
+     * Id of the access token
+     * @var string
+     */
+    protected $accessTokenId;
+
+    /**
+     * Set the ID of the associated access token
+     * @param  string $accessToken
+     * @return self
+     */
+    public function setAccessTokenId($accessTokenId)
+    {
+        $this->accessTokenId = $accessTokenId;
+
+        return $this;
+    }
 
     /**
      * Associate an access token
      * @param  \League\OAuth2\Server\Entity\AccessTokenEntity $accessToken
      * @return self
      */
-    public function setAccessToken(AccessTokenEntity $accessToken)
+    public function setAccessToken(AccessTokenEntity $accessTokenEntity)
     {
-        $this->accessToken = $accessToken;
+        $this->accessTokenEntity = $accessTokenEntity;
 
         return $this;
     }
@@ -40,11 +58,11 @@ class RefreshTokenEntity extends AbstractTokenEntity
      */
     public function getAccessToken()
     {
-        if (! $this->accessToken instanceof AccessTokenEntity) {
-            $this->accessToken = $this->server->getStorage('access_token')->getByRefreshToken($this);
+        if (! $this->accessTokenEntity instanceof AccessTokenEntity) {
+            $this->accessTokenEntity = $this->server->getStorage('access_token')->get($this->accessTokenId);
         }
 
-        return $this->accessToken;
+        return $this->accessTokenEntity;
     }
 
     /**
diff --git a/tests/unit/Entity/RefreshTokenEntityTest.php b/tests/unit/Entity/RefreshTokenEntityTest.php
index d0bcc162..6a57436d 100644
--- a/tests/unit/Entity/RefreshTokenEntityTest.php
+++ b/tests/unit/Entity/RefreshTokenEntityTest.php
@@ -8,8 +8,21 @@ use League\OAuth2\Server\Entity\AccessTokenEntity;
 use League\OAuth2\Server\Entity\RefreshTokenEntity;
 use \Mockery as M;
 
-class RefreshTokenTest extends \PHPUnit_Framework_TestCase
+class RefreshTokenEntityTest extends \PHPUnit_Framework_TestCase
 {
+    public function testSetAccessTokenId()
+    {
+        $server = M::mock('League\OAuth2\Server\AbstractServer');
+        $entity = new RefreshTokenEntity($server);
+        $entity->setAccessTokenId('foobar');
+
+        $reflector = new \ReflectionClass($entity);
+        $accessTokenProperty = $reflector->getProperty('accessTokenId');
+        $accessTokenProperty->setAccessible(true);
+
+        $this->assertSame($accessTokenProperty->getValue($entity), 'foobar');
+    }
+
     public function testSetAccessToken()
     {
         $server = M::mock('League\OAuth2\Server\AbstractServer');
@@ -17,7 +30,7 @@ class RefreshTokenTest extends \PHPUnit_Framework_TestCase
         $entity->setAccessToken((new AccessTokenEntity($server)));
 
         $reflector = new \ReflectionClass($entity);
-        $accessTokenProperty = $reflector->getProperty('accessToken');
+        $accessTokenProperty = $reflector->getProperty('accessTokenEntity');
         $accessTokenProperty->setAccessible(true);
 
         $this->assertTrue($accessTokenProperty->getValue($entity) instanceof AccessTokenEntity);
@@ -38,7 +51,7 @@ class RefreshTokenTest extends \PHPUnit_Framework_TestCase
 
         $accessTokenStorage = M::mock('League\OAuth2\Server\Storage\AccessTokenInterface');
         $accessTokenStorage->shouldReceive('setServer');
-        $accessTokenStorage->shouldReceive('getByRefreshToken')->andReturn(
+        $accessTokenStorage->shouldReceive('get')->andReturn(
             (new AccessTokenEntity($server))->setId('foobar')
         );
         $accessTokenStorage->shouldReceive('getScopes')->andReturn([