diff --git a/config/module.config.php b/config/module.config.php
index a056800..604820f 100644
--- a/config/module.config.php
+++ b/config/module.config.php
@@ -78,7 +78,7 @@
             'page_size' => 10,
             'page_size_param' => 'page_size',
             'collection_class' => 'Zend\Paginator\Paginator',
-            'entity_class' => 'Strapieno\UserAvatar\Model\Entity\UserAvatarEntity',
+            'entity_class' => 'Strapieno\UserAvatar\Model\Entity\AvatarEntity',
         ]
     ],
     'zf-content-negotiation' => [
@@ -98,7 +98,7 @@
      'zf-hal' => [
         // map each class (by name) to their metadata mappings
        'metadata_map' => [
-            'Strapieno\UserAvatar\Model\Entity\UserAvatarEntity' => [
+            'Strapieno\UserAvatar\Model\Entity\AvatarEntity' => [
                 'entity_identifier_name' => 'id',
                 'route_name' => 'api-rest/user/avatar',
                 'route_identifier_name' => 'user_id',
diff --git a/src/Listener/UserRestListener.php b/src/Listener/UserRestListener.php
index de488e8..f4172d3 100644
--- a/src/Listener/UserRestListener.php
+++ b/src/Listener/UserRestListener.php
@@ -55,7 +55,7 @@ public function onPostUpdate(ResourceEvent $e)
         $id  = $e->getParam('id');
         $user = $this->getUserFromId($id);
         $image = $e->getParam('image');
-
+        $image->setId($id);
         if ($user instanceof AvatarAwareInterface && $user instanceof ActiveRecordInterface) {
 
             $user->setAvatar($this->getUrlFromImage($image, $serviceLocator));
diff --git a/tests/Listener/UserRestListenerTest.php b/tests/Listener/UserRestListenerTest.php
index 82e515e..6220fe5 100644
--- a/tests/Listener/UserRestListenerTest.php
+++ b/tests/Listener/UserRestListenerTest.php
@@ -26,7 +26,7 @@ class UserRestListenerTest extends \PHPUnit_Framework_TestCase
                     'user' => [
                         'type' => 'Segment',
                         'options' => [
-                            'route' => '/user',
+                            'route' => '/user[/:user_id]',
                         ],
                         'child_routes' => [
                             'avatar' => [