diff --git a/Tests/Export/ExportManagerTest.php b/Tests/Export/ExportManagerTest.php
index 0b7f0cbee..dacf0a134 100644
--- a/Tests/Export/ExportManagerTest.php
+++ b/Tests/Export/ExportManagerTest.php
@@ -542,14 +542,15 @@ class ExportManagerTest extends KernelTestCase
         
         $export = $this->prophet->prophesize();
         $export->willImplement(ExportInterface::class);
+        $em = $this->container->get('doctrine.orm.entity_manager');
         $export->initiateQuery(
-                    Argument::Type(QueryBuilder::class), 
                     Argument::is(array('foo')), 
                     Argument::Type('array'),
                     Argument::is(array('a' => 'b'))
                 )
-                ->will(function ($qb) {
-                    return $qb[0]->addSelect('COUNT(user.id) as export')
+                ->will(function () use ($em) {
+                    $qb = $em->createQueryBuilder();
+                    return $qb->addSelect('COUNT(user.id) as export')
                             ->from('ChillMainBundle:User', 'user');
                     
                 });