diff --git a/src/DoctrineWriter.php b/src/DoctrineWriter.php
index a481ad6..b202a78 100644
--- a/src/DoctrineWriter.php
+++ b/src/DoctrineWriter.php
@@ -3,6 +3,7 @@
 namespace Port\Doctrine;
 
 use Port\Doctrine\Exception\UnsupportedDatabaseTypeException;
+use Port\Doctrine\LookupStrategy\PrimaryKeyLookupStrategy;
 use Port\Writer;
 use Doctrine\Common\Util\Inflector;
 use Doctrine\DBAL\Logging\SQLLogger;
@@ -26,20 +27,6 @@ class DoctrineWriter implements Writer, Writer\FlushableWriter
      */
     protected $objectManager;
 
-    /**
-     * Fully qualified model name
-     *
-     * @var string
-     */
-    protected $objectName;
-
-    /**
-     * Doctrine object repository
-     *
-     * @var ObjectRepository
-     */
-    protected $objectRepository;
-
     /**
      * @var ClassMetadata
      */
@@ -60,57 +47,50 @@ class DoctrineWriter implements Writer, Writer\FlushableWriter
     protected $truncate = true;
 
     /**
-     * List of fields used to lookup an object
-     *
-     * @var array
+     * @var LookupStrategy
      */
-    protected $lookupFields = [];
+    private $lookupStrategy;
 
-    /**
-     * Method used for looking up the item
-     *
-     * @var array
-     */
-    protected $lookupMethod;
+    public static function withLookupStrategy(
+        ObjectManager $objectManager,
+        LookupStrategy $lookupStrategy
+    ) {
+        return new self($objectManager, null, null, $lookupStrategy);
+    }
 
     /**
      * Constructor
      *
-     * @param ObjectManager $objectManager
-     * @param string        $objectName
-     * @param string|array  $index         Field or fields to find current entities by
-     * @param string        $lookupMethod  Method used for looking up the item
+     * @param ObjectManager  $objectManager
+     * @param string         $objectName
+     * @param string|array   $index          Field or fields to find current
+     *                                       entities by
+     * @param string         $lookupMethod   Method used for looking up the item
+     * @param LookupStrategy $lookupStrategy
+     *
+     * @throws UnsupportedDatabaseTypeException
      */
     public function __construct(
         ObjectManager $objectManager,
         $objectName,
         $index = null,
-        $lookupMethod = 'findOneBy'
+        $lookupMethod = 'findOneBy',
+        LookupStrategy $lookupStrategy = null
     ) {
         $this->ensureSupportedObjectManager($objectManager);
         $this->objectManager = $objectManager;
-        $this->objectRepository = $objectManager->getRepository($objectName);
         $this->objectMetadata = $objectManager->getClassMetadata($objectName);
-        //translate objectName in case a namespace alias is used
-        $this->objectName = $this->objectMetadata->getName();
-        if ($index) {
-            if (is_array($index)) {
-                $this->lookupFields = $index;
-            } else {
-                $this->lookupFields = [$index];
+
+        if ($objectManager !== null && $index !== null) {
+            $lookupStrategy = (new PrimaryKeyLookupStrategy($objectManager, $objectName))
+                ->withIndex($index);
+
+            if ($lookupMethod) {
+                $lookupStrategy = $lookupStrategy->withLookupMethod($lookupMethod);
             }
-        }
 
-        if (!method_exists($this->objectRepository, $lookupMethod)) {
-            throw new \InvalidArgumentException(
-                sprintf(
-                    'Repository %s has no method %s',
-                    get_class($this->objectRepository),
-                    $lookupMethod
-                )
-            );
+            $this->lookupStrategy = $lookupStrategy;
         }
-        $this->lookupMethod = [$this->objectRepository, $lookupMethod];
     }
 
     /**
@@ -189,7 +169,7 @@ public function writeItem(array $item)
     public function flush()
     {
         $this->objectManager->flush();
-        $this->objectManager->clear($this->objectName);
+        $this->objectManager->clear($this->objectMetadata->getName());
     }
 
     /**
@@ -284,7 +264,9 @@ protected function truncateTable()
             $query = $connection->getDatabasePlatform()->getTruncateTableSQL($tableName, true);
             $connection->executeQuery($query);
         } elseif ($this->objectManager instanceof \Doctrine\ODM\MongoDB\DocumentManager) {
-            $this->objectManager->getDocumentCollection($this->objectName)->remove(array());
+            $this->objectManager->getDocumentCollection(
+                $this->objectMetadata->getName()
+            )->remove([]);
         }
     }
 
@@ -320,27 +302,13 @@ protected function reEnableLogging()
      */
     protected function findOrCreateItem(array $item)
     {
-        $object = null;
         // If the table was not truncated to begin with, find current object
         // first
         if (!$this->truncate) {
-            if (!empty($this->lookupFields)) {
-                $lookupConditions = array();
-                foreach ($this->lookupFields as $fieldName) {
-                    $lookupConditions[$fieldName] = $item[$fieldName];
-                }
-
-                $object = call_user_func($this->lookupMethod, $lookupConditions);
-            } else {
-                $object = $this->objectRepository->find(current($item));
-            }
-        }
-
-        if (!$object) {
-            return $this->getNewInstance();
+            return $this->lookupStrategy->lookup($item);
         }
 
-        return $object;
+        return $this->getNewInstance();
     }
 
     protected function ensureSupportedObjectManager(ObjectManager $objectManager)
diff --git a/src/LookupStrategy.php b/src/LookupStrategy.php
new file mode 100644
index 0000000..4f61670
--- /dev/null
+++ b/src/LookupStrategy.php
@@ -0,0 +1,18 @@
+<?php
+
+namespace Port\Doctrine;
+
+/**
+ * Finds existing objects in the database.
+ */
+interface LookupStrategy
+{
+    /**
+     * Look up an item in the database.
+     *
+     * @param array $item
+     *
+     * @return mixed | null Null if no object was found.
+     */
+    public function lookup(array $item);
+}
diff --git a/src/LookupStrategy/PrimaryKeyLookupStrategy.php b/src/LookupStrategy/PrimaryKeyLookupStrategy.php
new file mode 100644
index 0000000..6b974d8
--- /dev/null
+++ b/src/LookupStrategy/PrimaryKeyLookupStrategy.php
@@ -0,0 +1,106 @@
+<?php
+
+namespace Port\Doctrine\LookupStrategy;
+
+use Doctrine\Common\Persistence\ObjectManager;
+use Doctrine\Common\Persistence\ObjectRepository;
+use Port\Doctrine\LookupStrategy;
+
+/**
+ * Default lookup strategy using primary keys.
+ */
+class PrimaryKeyLookupStrategy implements LookupStrategy
+{
+    /**
+     * @var ObjectManager
+     */
+    private $objectManager;
+
+    /**
+     * @var ObjectRepository
+     */
+    private $objectRepository;
+
+    /**
+     * @var array
+     */
+    private $lookupFields;
+
+    /**
+     * @var string
+     */
+    private $lookupMethod = 'findOneBy';
+
+    /**
+     * @param ObjectManager $objectManager
+     * @param string        $objectName    Fully qualified model name
+     */
+    public function __construct(ObjectManager $objectManager, $objectName)
+    {
+        $this->objectManager = $objectManager;
+        $this->objectRepository = $objectManager->getRepository($objectName);
+        $this->lookupFields = $objectManager->getClassMetadata($objectName)->getIdentifierFieldNames();
+    }
+
+    /**
+     * @param string $field Field to current current objects by.
+     * @return self
+     */
+    public function withLookupField($field)
+    {
+        return $this->withLookupFields([$field]);
+    }
+
+    /**
+     * Create lookup strategy with index
+     *
+     * @param array $fields Fields to find current objects by.
+     *
+     * @return self
+     */
+    public function withLookupFields(array $fields)
+    {
+        $new = clone $this;
+        $this->lookupFields = $fields;
+
+        return $new;
+    }
+
+    /**
+     * Doctrine repository method for finding objects.
+     *
+     * @param string $lookupMethod
+     *
+     * @return self
+     */
+    public function withLookupMethod($lookupMethod)
+    {
+        if (!method_exists($this->objectRepository, $lookupMethod)) {
+            throw new \InvalidArgumentException(
+                sprintf(
+                    'Repository %s has no method %s',
+                    get_class($this->objectRepository),
+                    $lookupMethod
+                )
+            );
+        }
+
+        $new = clone $this;
+        $new->lookupMethod = [$this->objectRepository, $lookupMethod];
+
+        return $new;
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function lookup(array $item)
+    {
+        $lookupConditions = array();
+        foreach ($this->lookupFields as $fieldName) {
+            $lookupConditions[$fieldName] = $item[$fieldName];
+        }
+
+        return call_user_func($this->lookupMethod, $lookupConditions);
+    }
+}
diff --git a/tests/DoctrineWriterTest.php b/tests/DoctrineWriterTest.php
index 9e5f639..e2c890d 100644
--- a/tests/DoctrineWriterTest.php
+++ b/tests/DoctrineWriterTest.php
@@ -2,7 +2,9 @@
 
 namespace Port\Doctrine\Tests;
 
+use Doctrine\ORM\EntityManager;
 use Port\Doctrine\DoctrineWriter;
+use Port\Doctrine\LookupStrategy\PrimaryKeyLookupStrategy;
 use Port\Doctrine\Tests\Fixtures\Entity\TestEntity;
 
 class DoctrineWriterTest extends \PHPUnit_Framework_TestCase
@@ -57,11 +59,7 @@ public function testUnsupportedDatabaseTypeException()
     protected function getEntityManager()
     {
         $em = $this->getMockBuilder('Doctrine\ORM\EntityManager')
-            ->setMethods(array('getRepository', 'getClassMetadata', 'persist', 'flush', 'clear', 'getConnection', 'getReference'))
-            ->disableOriginalConstructor()
-            ->getMock();
-
-        $repo = $this->getMockBuilder('Doctrine\ORM\EntityRepository')
+            ->setMethods(['getRepository', 'getClassMetadata', 'persist', 'flush', 'clear', 'getConnection', 'getReference'])
             ->disableOriginalConstructor()
             ->getMock();
 
@@ -112,10 +110,6 @@ protected function getEntityManager()
             ->method('executeQuery')
             ->with('TRUNCATE SQL');
 
-        $em->expects($this->once())
-            ->method('getRepository')
-            ->will($this->returnValue($repo));
-
         $em->expects($this->once())
             ->method('getClassMetadata')
             ->will($this->returnValue($metadata));
@@ -139,11 +133,7 @@ protected function getEntityManager()
     protected function getMongoDocumentManager()
     {
         $dm = $this->getMockBuilder('Doctrine\ODM\MongoDB\DocumentManager')
-            ->setMethods(array('getRepository', 'getClassMetadata', 'persist', 'flush', 'clear', 'getConnection', 'getDocumentCollection'))
-            ->disableOriginalConstructor()
-            ->getMock();
-
-        $repo = $this->getMockBuilder('Doctrine\ODM\MongoDB\DocumentRepository')
+            ->setMethods(['getClassMetadata', 'persist', 'flush', 'clear', 'getConnection', 'getDocumentCollection'])
             ->disableOriginalConstructor()
             ->getMock();
 
@@ -192,10 +182,6 @@ protected function getMongoDocumentManager()
         $connection->expects($this->never())
             ->method('executeQuery');
 
-        $dm->expects($this->once())
-            ->method('getRepository')
-            ->will($this->returnValue($repo));
-
         $dm->expects($this->once())
             ->method('getClassMetadata')
             ->will($this->returnValue($metadata));
@@ -284,4 +270,32 @@ public function testFlushAndClear()
         $writer = new DoctrineWriter($em, 'Port:TestEntity');
         $writer->finish();
     }
+
+    public function testWithLookupStrategy()
+    {
+        $em = $this->getEntityManager();
+//        $em = $this->getMockBuilder(EntityManager::class)
+//            ->disableOriginalConstructor()
+//            ->getMock();
+
+//        $repo = $this->getMockBuilder('Doctrine\ORM\EntityRepository')
+//            ->disableOriginalConstructor()
+//            ->getMock();
+//
+//        $em->expects($this->once())
+//            ->method('getRepository')
+//            ->with($this->equalTo('Port:TestEntity'))
+//            ->willReturn($repo);
+
+        $lookupStrategy = new PrimaryKeyLookupStrategy($em, 'Port:TestEntity');
+        $writer = DoctrineWriter::withLookupStrategy($this->getEntityManager(), $lookupStrategy);
+
+        $item = [
+            'firstProperty' => 'some value',
+            'secondProperty' => 'some other value',
+            'firstAssociation'=> new TestEntity(),
+        ];
+        $writer->writeItem($item);
+        $writer->finish();
+    }
 }