diff --git a/src/Shopify.php b/src/Shopify.php
index db6e113fdc805d293d3c5d32ad8f493c4b42cd4e..7647816933e0a7596929cc40f17b7f8f0027f916 100644
--- a/src/Shopify.php
+++ b/src/Shopify.php
@@ -4,14 +4,15 @@ declare(strict_types=1);
 
 namespace Octolize\Shopify\Codeception;
 
+use Codeception\Module\WebDriver;
 use Exception;
 
-class Shopify extends \Codeception\Module
+class Shopify extends WebDriver
 {
 
     public function loginAsAdmin()
     {
-        $i = $this->getModule('WebDriver');
+        $i    = $this;
         $isPl = false;
         if ($i->loadSessionSnapshot('admin')) {
             return;
@@ -28,7 +29,7 @@ class Shopify extends \Codeception\Module
                 $i->waitForText('Kontynuuj z Google', 30);
                 $isPl = true;
             }
-            $i->click($isPl?'Kontynuuj z Google':'Continue with Google');
+            $i->click($isPl ? 'Kontynuuj z Google' : 'Continue with Google');
             try {
                 $i->seeInCurrentUrl(parse_url($this->config['admin_url'], PHP_URL_PATH));
                 $i->dontSeeInCurrentUrl('google.com');
@@ -58,6 +59,7 @@ class Shopify extends \Codeception\Module
                 }
             }
         }
+        $i->wait(rand(1, 5));
         $i->saveSessionSnapshot('admin');
     }