diff --git a/src/Bankiru/Api/Doctrine/EntityRepository.php b/src/Bankiru/Api/Doctrine/EntityRepository.php
index 380155e..51c5644 100644
--- a/src/Bankiru/Api/Doctrine/EntityRepository.php
+++ b/src/Bankiru/Api/Doctrine/EntityRepository.php
@@ -4,6 +4,7 @@
 
 use Bankiru\Api\Doctrine\Mapping\ApiMetadata;
 use Bankiru\Api\Doctrine\Proxy\ApiCollection;
+use Doctrine\Common\Collections\Collection;
 use Doctrine\Common\Persistence\ObjectRepository;
 use ScayTrase\Api\Rpc\RpcClientInterface;
 
@@ -76,7 +77,7 @@ public function findAll()
      */
     public function findBy(array $criteria, array $orderBy = null, $limit = null, $offset = null)
     {
-        return new ApiCollection($this->manager, $this->metadata, [$criteria, $orderBy, $limit, $offset]);
+        return $this->createLazyCollection($criteria, $orderBy, $limit, $offset)->toArray();
     }
 
     /**
@@ -101,6 +102,19 @@ public function getManager()
         return $this->manager;
     }
 
+    /**
+     * @param array      $criteria
+     * @param array|null $orderBy
+     * @param int|null   $limit
+     * @param int|null   $offset
+     *
+     * @return Collection
+     */
+    public function createLazyCollection(array $criteria, array $orderBy = null, $limit = null, $offset = null)
+    {
+        return new ApiCollection($this->manager, $this->metadata, [$criteria, $orderBy, $limit, $offset]);
+    }
+
     /**
      * @return RpcClientInterface
      */
@@ -116,4 +130,16 @@ protected function getMetadata()
     {
         return $this->metadata;
     }
+
+    /**
+     * Hydrates object from given data or merges it to already fetched object
+     *
+     * @param mixed $data
+     *
+     * @return object
+     */
+    protected function hydrateObject($data)
+    {
+        return $this->getManager()->getUnitOfWork()->getOrCreateEntity($this->getClassName(), $data);
+    }
 }
diff --git a/src/Bankiru/Api/Doctrine/Proxy/ApiCollection.php b/src/Bankiru/Api/Doctrine/Proxy/ApiCollection.php
index 7b52138..3949725 100644
--- a/src/Bankiru/Api/Doctrine/Proxy/ApiCollection.php
+++ b/src/Bankiru/Api/Doctrine/Proxy/ApiCollection.php
@@ -2,7 +2,7 @@
 
 namespace Bankiru\Api\Doctrine\Proxy;
 
-use Bankiru\Api\Doctrine\EntityManager;
+use Bankiru\Api\Doctrine\ApiEntityManager;
 use Bankiru\Api\Doctrine\Mapping\ApiMetadata;
 use Doctrine\Common\Collections\AbstractLazyCollection;
 use Doctrine\Common\Collections\ArrayCollection;
@@ -11,7 +11,7 @@
 
 class ApiCollection extends AbstractLazyCollection
 {
-    /** @var  EntityManager */
+    /** @var  ApiEntityManager */
     private $manager;
     /** @var ApiMetadata */
     private $metadata;
@@ -25,13 +25,13 @@ class ApiCollection extends AbstractLazyCollection
     /**
      * ApiCollection constructor.
      *
-     * @param EntityManager $manager
-     * @param ApiMetadata   $class
-     * @param array         $searchArgs
-     * @param Collection    $collection
+     * @param ApiEntityManager $manager
+     * @param ApiMetadata      $class
+     * @param array            $searchArgs
+     * @param Collection       $collection
      */
     public function __construct(
-        EntityManager $manager,
+        ApiEntityManager $manager,
         ApiMetadata $class,
         array $searchArgs,
         Collection $collection = null
@@ -46,7 +46,7 @@ public function __construct(
 
     /**
      * @param object $owner
-     * @param        $assoc
+     * @param array  $assoc
      */
     public function setOwner($owner, array $assoc)
     {
diff --git a/src/Bankiru/Api/Tests/CollectionLoadingTest.php b/src/Bankiru/Api/Tests/CollectionLoadingTest.php
index 07bb626..ca93e04 100644
--- a/src/Bankiru/Api/Tests/CollectionLoadingTest.php
+++ b/src/Bankiru/Api/Tests/CollectionLoadingTest.php
@@ -2,6 +2,7 @@
 
 namespace Bankiru\Api\Tests;
 
+use Bankiru\Api\Doctrine\EntityRepository;
 use Bankiru\Api\Doctrine\Proxy\ApiCollection;
 use Bankiru\Api\Test\Entity\Sub\SubEntity;
 use Doctrine\Common\Collections\ArrayCollection;
@@ -9,22 +10,19 @@
 
 class CollectionLoadingTest extends AbstractEntityManagerTest
 {
-    protected function getClientNames()
-    {
-        return [self::DEFAULT_CLIENT, 'test-reference-client'];
-    }
-
     public function testLazyCollections()
     {
+        /** @var EntityRepository $repository */
         $repository = $this->getManager()->getRepository(SubEntity::class);
-        /** @var SubEntity[]|ArrayCollection|ApiCollection $entities */
-        $entities = $repository->findBy(['subPayload' => 'sub-payload']);
+        /** @var ApiCollection $collection */
+        $collection = $repository->createLazyCollection(['subPayload' => 'sub-payload']);
 
-        self::assertInstanceOf(ApiCollection::class, $entities);
-        self::assertFalse($entities->isInitialized());
+        self::assertInstanceOf(ApiCollection::class, $collection);
+        self::assertFalse($collection->isInitialized());
 
         try {
-            $entities->count();
+            $collection->count();
+            self::fail('Should fail');
         } catch (\OutOfBoundsException $exception) {
             self::assertEquals('Mock queue is empty', $exception->getMessage());
         }
@@ -69,7 +67,7 @@ public function testFindBy()
 
         /** @var SubEntity[]|ArrayCollection $entities */
         $entities = $repository->findBy(['subPayload' => 'sub-payload']);
-        self::assertInstanceOf(\Countable::class, $entities);
+        self::assertInternalType('array', $entities);
         self::assertCount(3, $entities);
 
         foreach ($entities as $entity) {
@@ -83,4 +81,9 @@ public function testFindBy()
             }
         }
     }
+
+    protected function getClientNames()
+    {
+        return [self::DEFAULT_CLIENT, 'test-reference-client'];
+    }
 }
diff --git a/src/Bankiru/Api/Tests/ReferenceLoadingTest.php b/src/Bankiru/Api/Tests/ReferenceLoadingTest.php
index 4238059..1f5e762 100644
--- a/src/Bankiru/Api/Tests/ReferenceLoadingTest.php
+++ b/src/Bankiru/Api/Tests/ReferenceLoadingTest.php
@@ -191,8 +191,7 @@ public function testFindByReference()
         self::assertCount(1, $children);
         /** @var TestEntity $child */
 
-        $childrenArray = $children->toArray();
-        $child         = array_shift($childrenArray);
+        $child         = array_shift($children);
 
         self::assertEquals($parent, $child->getParent());
         self::assertEquals($parent->getPayload(), $child->getParent()->getPayload());