diff --git a/Tests/Unit/DependencyInjection/DoctrineEncryptExtensionTest.php b/Tests/Unit/DependencyInjection/DoctrineEncryptExtensionTest.php
index 42b422c8..e2f792b6 100644
--- a/Tests/Unit/DependencyInjection/DoctrineEncryptExtensionTest.php
+++ b/Tests/Unit/DependencyInjection/DoctrineEncryptExtensionTest.php
@@ -26,7 +26,7 @@ public function testConfigLoadHalite()
         $container = $this->createContainer();
         $this->extension->load([[]], $container);
 
-        $this->assertSame(HaliteEncryptor::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
+        self::assertSame(HaliteEncryptor::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
     }
 
     public function testConfigLoadDefuse()
@@ -38,7 +38,7 @@ public function testConfigLoadDefuse()
         ];
         $this->extension->load([$config], $container);
 
-        $this->assertSame(DefuseEncryptor::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
+        self::assertSame(DefuseEncryptor::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
     }
 
     public function testConfigLoadCustom()
@@ -49,17 +49,15 @@ public function testConfigLoadCustom()
         ];
         $this->extension->load([$config], $container);
 
-        $this->markTestSkipped();
+        self::markTestSkipped();
 
-        $this->assertSame(self::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
+        self::assertSame(self::class, $container->getParameter('ambta_doctrine_encrypt.encryptor_class_name'));
     }
 
     private function createContainer()
     {
-        $container = new ContainerBuilder(
+        return new ContainerBuilder(
             new ParameterBag(['kernel.debug' => false])
         );
-
-        return $container;
     }
 }
diff --git a/Tests/Unit/Subscribers/DoctrineEncryptSubscriberTest.php b/Tests/Unit/Subscribers/DoctrineEncryptSubscriberTest.php
index 59b880cb..7ee6189c 100644
--- a/Tests/Unit/Subscribers/DoctrineEncryptSubscriberTest.php
+++ b/Tests/Unit/Subscribers/DoctrineEncryptSubscriberTest.php
@@ -30,31 +30,26 @@ class DoctrineEncryptSubscriberTest extends TestCase
      */
     private $encryptor;
 
-    /**
-     * @var Reader|MockObject
-     */
-    private $reader;
-
     protected function setUp()
     {
         $this->encryptor = $this->createMock(EncryptorInterface::class);
         $this->encryptor
-            ->expects($this->any())
+            ->expects(self::any())
             ->method('encrypt')
             ->willReturnCallback(function (string $arg) {
                 return 'encrypted-'.$arg;
             })
         ;
         $this->encryptor
-            ->expects($this->any())
+            ->expects(self::any())
             ->method('decrypt')
             ->willReturnCallback(function (string $arg) {
                 return preg_replace('/^encrypted-/', '', $arg);
             })
         ;
 
-        $this->reader = $this->createMock(Reader::class);
-        $this->reader->expects($this->any())
+        $reader = $this->createMock(Reader::class);
+        $reader->expects(self::any())
             ->method('getPropertyAnnotation')
             ->willReturnCallback(function (\ReflectionProperty $reflProperty, string $class) {
                 if (Encrypted::class === $class) {
@@ -68,115 +63,133 @@ protected function setUp()
             })
         ;
 
-        $this->subscriber = new DoctrineEncryptSubscriber($this->reader, $this->encryptor);
+        $this->subscriber = new DoctrineEncryptSubscriber($reader, $this->encryptor);
     }
 
-    public function testSetRestorEncryptor()
+    public function testSetRestoreEncryptor()
     {
         $replaceEncryptor = $this->createMock(EncryptorInterface::class);
 
-        $this->assertSame($this->encryptor, $this->subscriber->getEncryptor());
+        self::assertSame($this->encryptor, $this->subscriber->getEncryptor());
         $this->subscriber->setEncryptor($replaceEncryptor);
-        $this->assertSame($replaceEncryptor, $this->subscriber->getEncryptor());
+        self::assertSame($replaceEncryptor, $this->subscriber->getEncryptor());
         $this->subscriber->restoreEncryptor();
-        $this->assertSame($this->encryptor, $this->subscriber->getEncryptor());
+        self::assertSame($this->encryptor, $this->subscriber->getEncryptor());
     }
 
     public function testProcessFieldsEncrypt()
     {
         $user = new User('David', 'Switzerland');
 
-        $this->subscriber->processFields($user, true);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, true);
 
-        $this->assertStringStartsWith('encrypted-', $user->name);
-        $this->assertStringStartsWith('encrypted-', $user->getAddress());
+        self::assertStringStartsWith('encrypted-', $user->name);
+        self::assertStringStartsWith('encrypted-', $user->getAddress());
     }
 
     public function testProcessFieldsEncryptExtend()
     {
         $user = new ExtendedUser('David', 'Switzerland', 'extra');
 
-        $this->subscriber->processFields($user, true);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, true);
 
-        $this->assertStringStartsWith('encrypted-', $user->name);
-        $this->assertStringStartsWith('encrypted-', $user->getAddress());
-        $this->assertStringStartsWith('encrypted-', $user->extra);
+        self::assertStringStartsWith('encrypted-', $user->name);
+        self::assertStringStartsWith('encrypted-', $user->getAddress());
+        self::assertStringStartsWith('encrypted-', $user->extra);
     }
 
     public function testProcessFieldsEncryptEmbedded()
     {
         $withUser = new WithUser('Thing', 'foo', new User('David', 'Switzerland'));
 
-        $this->subscriber->processFields($withUser, true);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($withUser, $em, true);
 
-        $this->assertStringStartsWith('encrypted-', $withUser->name);
-        $this->assertSame('foo', $withUser->foo);
-        $this->assertStringStartsWith('encrypted-', $withUser->user->name);
-        $this->assertStringStartsWith('encrypted-', $withUser->user->getAddress());
+        self::assertStringStartsWith('encrypted-', $withUser->name);
+        self::assertSame('foo', $withUser->foo);
+        self::assertStringStartsWith('encrypted-', $withUser->user->name);
+        self::assertStringStartsWith('encrypted-', $withUser->user->getAddress());
     }
 
     public function testProcessFieldsEncryptNull()
     {
         $user = new User('David', null);
 
-        $this->subscriber->processFields($user, true);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, true);
 
-        $this->assertStringStartsWith('encrypted-', $user->name);
-        $this->assertNull($user->getAddress());
+        self::assertStringStartsWith('encrypted-', $user->name);
+        self::assertNull($user->getAddress());
     }
 
     public function testProcessFieldsNoEncryptor()
     {
         $user = new User('David', 'Switzerland');
 
+        $em = $this->createMock(EntityManagerInterface::class);
+
         $this->subscriber->setEncryptor(null);
-        $this->subscriber->processFields($user, true);
+        $this->subscriber->processFields($user, $em, true);
 
-        $this->assertSame('David', $user->name);
-        $this->assertSame('Switzerland', $user->getAddress());
+        self::assertSame('David', $user->name);
+        self::assertSame('Switzerland', $user->getAddress());
     }
 
     public function testProcessFieldsDecrypt()
     {
         $user = new User('encrypted-David<ENC>', 'encrypted-Switzerland<ENC>');
 
-        $this->subscriber->processFields($user, false);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, false);
 
-        $this->assertSame('David', $user->name);
-        $this->assertSame('Switzerland', $user->getAddress());
+        self::assertSame('David', $user->name);
+        self::assertSame('Switzerland', $user->getAddress());
     }
 
     public function testProcessFieldsDecryptExtended()
     {
         $user = new ExtendedUser('encrypted-David<ENC>', 'encrypted-Switzerland<ENC>', 'encrypted-extra<ENC>');
 
-        $this->subscriber->processFields($user, false);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, false);
 
-        $this->assertSame('David', $user->name);
-        $this->assertSame('Switzerland', $user->getAddress());
-        $this->assertSame('extra', $user->extra);
+        self::assertSame('David', $user->name);
+        self::assertSame('Switzerland', $user->getAddress());
+        self::assertSame('extra', $user->extra);
     }
 
     public function testProcessFieldsDecryptEmbedded()
     {
         $withUser = new WithUser('encrypted-Thing<ENC>', 'foo', new User('encrypted-David<ENC>', 'encrypted-Switzerland<ENC>'));
 
-        $this->subscriber->processFields($withUser, false);
+        $em = $this->createMock(EntityManagerInterface::class);
 
-        $this->assertSame('Thing', $withUser->name);
-        $this->assertSame('foo', $withUser->foo);
-        $this->assertSame('David', $withUser->user->name);
-        $this->assertSame('Switzerland', $withUser->user->getAddress());
+        $this->subscriber->processFields($withUser, $em, false);
+
+        self::assertSame('Thing', $withUser->name);
+        self::assertSame('foo', $withUser->foo);
+        self::assertSame('David', $withUser->user->name);
+        self::assertSame('Switzerland', $withUser->user->getAddress());
     }
 
     public function testProcessFieldsDecryptNull()
     {
         $user = new User('encrypted-David<ENC>', null);
 
-        $this->subscriber->processFields($user, false);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, false);
 
-        $this->assertSame('David', $user->name);
-        $this->assertNull($user->getAddress());
+        self::assertSame('David', $user->name);
+        self::assertNull($user->getAddress());
     }
 
     public function testProcessFieldsDecryptNonEncrypted()
@@ -184,10 +197,12 @@ public function testProcessFieldsDecryptNonEncrypted()
         // no trailing <ENC> but somethint that our mock decrypt would change if called
         $user = new User('encrypted-David', 'encrypted-Switzerland');
 
-        $this->subscriber->processFields($user, false);
+        $em = $this->createMock(EntityManagerInterface::class);
+
+        $this->subscriber->processFields($user, $em, false);
 
-        $this->assertSame('encrypted-David', $user->name);
-        $this->assertSame('encrypted-Switzerland', $user->getAddress());
+        self::assertSame('encrypted-David', $user->name);
+        self::assertSame('encrypted-Switzerland', $user->getAddress());
     }
 
     /**
@@ -198,25 +213,25 @@ public function testOnFlush()
         $user = new User('David', 'Switzerland');
 
         $uow = $this->createMock(UnitOfWork::class);
-        $uow->expects($this->any())
+        $uow->expects(self::any())
             ->method('getScheduledEntityInsertions')
             ->willReturn([$user])
         ;
         $em = $this->createMock(EntityManagerInterface::class);
-        $em->expects($this->any())
+        $em->expects(self::any())
             ->method('getUnitOfWork')
             ->willReturn($uow)
         ;
         $classMetaData = $this->createMock(ClassMetadata::class);
-        $em->expects($this->once())->method('getClassMetadata')->willReturn($classMetaData);
-        $uow->expects($this->once())->method('recomputeSingleEntityChangeSet');
+        $em->expects(self::any())->method('getClassMetadata')->willReturn($classMetaData);
+        $uow->expects(self::any())->method('recomputeSingleEntityChangeSet');
 
         $onFlush = new OnFlushEventArgs($em);
 
         $this->subscriber->onFlush($onFlush);
 
-        $this->assertStringStartsWith('encrypted-', $user->name);
-        $this->assertStringStartsWith('encrypted-', $user->getAddress());
+        self::assertStringStartsWith('encrypted-', $user->name);
+        self::assertStringStartsWith('encrypted-', $user->getAddress());
     }
 
     /**
@@ -227,12 +242,12 @@ public function testPostFlush()
         $user = new User('encrypted-David<ENC>', 'encrypted-Switzerland<ENC>');
 
         $uow = $this->createMock(UnitOfWork::class);
-        $uow->expects($this->any())
+        $uow->expects(self::any())
             ->method('getIdentityMap')
             ->willReturn([[$user]])
         ;
         $em = $this->createMock(EntityManagerInterface::class);
-        $em->expects($this->any())
+        $em->expects(self::any())
             ->method('getUnitOfWork')
             ->willReturn($uow)
         ;
@@ -240,7 +255,7 @@ public function testPostFlush()
 
         $this->subscriber->postFlush($postFlush);
 
-        $this->assertSame('David', $user->name);
-        $this->assertSame('Switzerland', $user->getAddress());
+        self::assertSame('David', $user->name);
+        self::assertSame('Switzerland', $user->getAddress());
     }
 }
diff --git a/src/Command/DoctrineEncryptDatabaseCommand.php b/src/Command/DoctrineEncryptDatabaseCommand.php
index f5d371c8..3decb3ac 100644
--- a/src/Command/DoctrineEncryptDatabaseCommand.php
+++ b/src/Command/DoctrineEncryptDatabaseCommand.php
@@ -84,7 +84,7 @@ protected function execute(InputInterface $input, OutputInterface $output)
             $output->writeln(sprintf('Processing <comment>%s</comment>', $metaData->name));
             $progressBar = new ProgressBar($output, $totalCount);
             foreach ($iterator as $row) {
-                $this->subscriber->processFields($row[0]);
+                $this->subscriber->processFields($row[0], $this->entityManager);
 
                 if (($i % $batchSize) === 0) {
                     $this->entityManager->flush();
diff --git a/src/Subscribers/DoctrineEncryptSubscriber.php b/src/Subscribers/DoctrineEncryptSubscriber.php
index 85207851..2ed3ba26 100644
--- a/src/Subscribers/DoctrineEncryptSubscriber.php
+++ b/src/Subscribers/DoctrineEncryptSubscriber.php
@@ -2,6 +2,7 @@
 
 namespace Ambta\DoctrineEncryptBundle\Subscribers;
 
+use Doctrine\ORM\EntityManagerInterface;
 use Doctrine\ORM\Event\PreFlushEventArgs;
 use ReflectionClass;
 use Doctrine\ORM\Event\PostFlushEventArgs;
@@ -44,7 +45,7 @@ class DoctrineEncryptSubscriber implements EventSubscriber
 
     /**
      * Annotation reader
-     * @var \Doctrine\Common\Annotations\Reader
+     * @var Reader
      */
     private $annReader;
 
@@ -85,7 +86,7 @@ public function __construct(Reader $annReader, EncryptorInterface $encryptor)
     /**
      * Change the encryptor
      *
-     * @param EncryptorInterface $encryptor
+     * @param EncryptorInterface|null $encryptor
      */
     public function setEncryptor(EncryptorInterface $encryptor = null)
     {
@@ -117,36 +118,36 @@ public function restoreEncryptor()
      * So for example after form submit the preUpdate encrypted the entity
      * We have to decrypt them before showing them again.
      *
-     * @param LifecycleEventArgs $args
+     * @param LifecycleEventArgs $liveCycleEventArgs
      */
-    public function postUpdate(LifecycleEventArgs $args)
+    public function postUpdate(LifecycleEventArgs $liveCycleEventArgs)
     {
-        $entity = $args->getEntity();
-        $this->processFields($entity, false);
+        $entity = $liveCycleEventArgs->getEntity();
+        $this->processFields($entity, $liveCycleEventArgs->getEntityManager(), false);
     }
 
     /**
      * Listen a preUpdate lifecycle event.
      * Encrypt entities property's values on preUpdate, so they will be stored encrypted
      *
-     * @param PreUpdateEventArgs $args
+     * @param PreUpdateEventArgs $preUpdateEventArgs
      */
-    public function preUpdate(PreUpdateEventArgs $args)
+    public function preUpdate(PreUpdateEventArgs $preUpdateEventArgs)
     {
-        $entity = $args->getEntity();
-        $this->processFields($entity);
+        $entity = $preUpdateEventArgs->getEntity();
+        $this->processFields($entity, $preUpdateEventArgs->getEntityManager());
     }
 
     /**
      * Listen a postLoad lifecycle event.
      * Decrypt entities property's values when loaded into the entity manger
      *
-     * @param LifecycleEventArgs $args
+     * @param LifecycleEventArgs $liveCycleEventArgs
      */
-    public function postLoad(LifecycleEventArgs $args)
+    public function postLoad(LifecycleEventArgs $liveCycleEventArgs)
     {
-        $entity = $args->getEntity();
-        $this->processFields($entity, false);
+        $entity = $liveCycleEventArgs->getEntity();
+        $this->processFields($entity, $liveCycleEventArgs->getEntityManager(), false);
     }
 
     /**
@@ -161,7 +162,7 @@ public function preFlush(PreFlushEventArgs $preFlushEventArgs)
         foreach ($unitOfWOrk->getIdentityMap() as $entityName => $entityArray) {
             if (isset($this->cachedDecryptions[$entityName])) {
                 foreach ($entityArray as $entityId => $instance) {
-                    $this->processFields($instance);
+                    $this->processFields($instance, $preFlushEventArgs->getEntityManager());
                 }
             }
         }
@@ -179,7 +180,7 @@ public function onFlush(OnFlushEventArgs $onFlushEventArgs)
         $unitOfWork = $onFlushEventArgs->getEntityManager()->getUnitOfWork();
         foreach ($unitOfWork->getScheduledEntityInsertions() as $entity) {
             $encryptCounterBefore = $this->encryptCounter;
-            $this->processFields($entity);
+            $this->processFields($entity, $onFlushEventArgs->getEntityManager());
             if ($this->encryptCounter > $encryptCounterBefore ) {
                 $classMetadata = $onFlushEventArgs->getEntityManager()->getClassMetadata(get_class($entity));
                 $unitOfWork->recomputeSingleEntityChangeSet($classMetadata, $entity);
@@ -198,7 +199,7 @@ public function postFlush(PostFlushEventArgs $postFlushEventArgs)
         $unitOfWork = $postFlushEventArgs->getEntityManager()->getUnitOfWork();
         foreach ($unitOfWork->getIdentityMap() as $entityMap) {
             foreach ($entityMap as $entity) {
-                $this->processFields($entity, false);
+                $this->processFields($entity, $postFlushEventArgs->getEntityManager(), false);
             }
         }
     }
@@ -223,14 +224,13 @@ public function getSubscribedEvents()
     /**
      * Process (encrypt/decrypt) entities fields
      *
-     * @param Object $entity doctrine entity
+     * @param Object  $entity             doctrine entity
+     * @param EntityManagerInterface $em  entity manager
      * @param Boolean $isEncryptOperation If true - encrypt, false - decrypt entity
      *
-     * @throws \RuntimeException
-     *
      * @return object|null
      */
-    public function processFields($entity, $isEncryptOperation = true)
+    public function processFields($entity, $em, $isEncryptOperation = true)
     {
         if (!empty($this->encryptor)) {
             // Check which operation to be used
@@ -238,13 +238,20 @@ public function processFields($entity, $isEncryptOperation = true)
 
             $realClass = ClassUtils::getClass($entity);
 
+            $className = get_class($entity);
+
+            $classMetadata = $em->getClassMetadata($className);
+            if(null !== $classMetadata) {
+                $className = $classMetadata->getName();
+            }
+
             // Get ReflectionClass of our entity
             $properties = $this->getClassProperties($realClass);
 
             // Foreach property in the reflection class
             foreach ($properties as $refProperty) {
                 if ($this->annReader->getPropertyAnnotation($refProperty, 'Doctrine\ORM\Mapping\Embedded')) {
-                    $this->handleEmbeddedAnnotation($entity, $refProperty, $isEncryptOperation);
+                    $this->handleEmbeddedAnnotation($entity, $refProperty, $em, $isEncryptOperation);
                     continue;
                 }
 
@@ -254,20 +261,18 @@ public function processFields($entity, $isEncryptOperation = true)
                 if ($this->annReader->getPropertyAnnotation($refProperty, self::ENCRYPTED_ANN_NAME)) {
                     $pac = PropertyAccess::createPropertyAccessor();
                     $value = $pac->getValue($entity, $refProperty->getName());
-                    if ($encryptorMethod == 'decrypt') {
-                        if (!is_null($value) and !empty($value)) {
-                            if (substr($value, -strlen(self::ENCRYPTION_MARKER)) == self::ENCRYPTION_MARKER) {
-                                $this->decryptCounter++;
-                                $currentPropValue = $this->encryptor->decrypt(substr($value, 0, -5));
-                                $pac->setValue($entity, $refProperty->getName(), $currentPropValue);
-                                $this->cachedDecryptions[get_class($entity)][spl_object_id($entity)][$refProperty->getName()][$currentPropValue] = $value;
-                            }
+                    if ($encryptorMethod === 'decrypt') {
+                        if (!empty($value) && substr($value, -strlen(self::ENCRYPTION_MARKER)) === self::ENCRYPTION_MARKER) {
+                            $this->decryptCounter++;
+                            $currentPropValue = $this->encryptor->decrypt(substr($value, 0, -strlen(self::ENCRYPTION_MARKER)));
+                            $pac->setValue($entity, $refProperty->getName(), $currentPropValue);
+                            $this->cachedDecryptions[$className][spl_object_id($entity)][$refProperty->getName()][$currentPropValue] = $value;
                         }
                     } else {
-                        if (!is_null($value) and !empty($value)) {
-                            if (isset($this->cachedDecryptions[get_class($entity)][spl_object_id($entity)][$refProperty->getName()][$value])) {
-                                $pac->setValue($entity, $refProperty->getName(), $this->cachedDecryptions[get_class($entity)][spl_object_id($entity)][$refProperty->getName()][$value]);
-                            } elseif (substr($value, -strlen(self::ENCRYPTION_MARKER)) != self::ENCRYPTION_MARKER) {
+                        if (!empty($value)) {
+                            if (isset($this->cachedDecryptions[$className][spl_object_id($entity)][$refProperty->getName()][$value])) {
+                                $pac->setValue($entity, $refProperty->getName(), $this->cachedDecryptions[$className][spl_object_id($entity)][$refProperty->getName()][$value]);
+                            } elseif (substr($value, -strlen(self::ENCRYPTION_MARKER)) !== self::ENCRYPTION_MARKER) {
                                 $this->encryptCounter++;
                                 $currentPropValue = $this->encryptor->encrypt($value).self::ENCRYPTION_MARKER;
                                 $pac->setValue($entity, $refProperty->getName(), $currentPropValue);
@@ -283,7 +288,13 @@ public function processFields($entity, $isEncryptOperation = true)
         return $entity;
     }
 
-    private function handleEmbeddedAnnotation($entity, ReflectionProperty $embeddedProperty, bool $isEncryptOperation = true)
+    /**
+     * @param                    $entity
+     * @param ReflectionProperty $embeddedProperty
+     * @param                    $em
+     * @param bool               $isEncryptOperation
+     */
+    private function handleEmbeddedAnnotation($entity, ReflectionProperty $embeddedProperty, $em, bool $isEncryptOperation = true)
     {
         $propName = $embeddedProperty->getName();
 
@@ -292,7 +303,7 @@ private function handleEmbeddedAnnotation($entity, ReflectionProperty $embeddedP
         $embeddedEntity = $pac->getValue($entity, $propName);
 
         if ($embeddedEntity) {
-            $this->processFields($embeddedEntity, $isEncryptOperation);
+            $this->processFields($embeddedEntity, $em, $isEncryptOperation);
         }
     }