diff --git a/test/e2e/driver/aut.py b/test/e2e/driver/aut.py
index 83c441a298..20147b0b0e 100644
--- a/test/e2e/driver/aut.py
+++ b/test/e2e/driver/aut.py
@@ -72,7 +72,7 @@ class AUT:
     def detach_context(self):
         if self.ctx is None:
             return
-        squish.currentApplicationContext().detach()
+        driver.currentApplicationContext().detach()
         self.ctx = None
 
     @allure.step('Attach Squish to Test Application')
@@ -83,7 +83,7 @@ class AUT:
             SquishServer().add_attachable_aut(self.aut_id, self.port)
             if self.ctx is None:
                 self.ctx = context.get_context(self.aut_id)
-            squish.setApplicationContext(self.ctx)
+            driver.setApplicationContext(self.ctx)
             assert squish.waitFor(lambda: self.ctx.isRunning, configs.timeouts.PROCESS_TIMEOUT_SEC)
         except Exception as err:
             LOG.error('Failed to attach AUT: %s', err)
diff --git a/test/e2e/driver/context.py b/test/e2e/driver/context.py
index 98c3b0b279..48f9f59fe8 100644
--- a/test/e2e/driver/context.py
+++ b/test/e2e/driver/context.py
@@ -4,6 +4,7 @@ import allure
 import squish
 
 import configs
+import driver
 from driver.server import SquishServer
 
 LOG = logging.getLogger(__name__)
@@ -13,7 +14,7 @@ LOG = logging.getLogger(__name__)
 def get_context(aut_id: str):
     LOG.info('Attaching to: %s', aut_id)
     try:
-        context = squish.attachToApplication(aut_id, SquishServer().host, SquishServer().port)
+        context = driver.attachToApplication(aut_id, SquishServer().host, SquishServer().port)
         if context is not None:
             LOG.info('AUT %s context found', aut_id)
             return context
@@ -24,7 +25,7 @@ def get_context(aut_id: str):
 
 @allure.step('Detaching')
 def detach():
-    for ctx in squish.applicationContextList():
+    for ctx in driver.applicationContextList():
         ctx.detach()
         assert squish.waitFor(lambda: not ctx.isRunning, configs.timeouts.APP_LOAD_TIMEOUT_MSEC)
     LOG.info('All AUTs detached')
diff --git a/test/e2e/gui/components/onboarding/beta_consent_popup.py b/test/e2e/gui/components/onboarding/beta_consent_popup.py
index e4fc5e45a6..4796969aea 100644
--- a/test/e2e/gui/components/onboarding/beta_consent_popup.py
+++ b/test/e2e/gui/components/onboarding/beta_consent_popup.py
@@ -19,5 +19,5 @@ class BetaConsentPopup(QObject):
     def confirm(self):
         self._agree_to_use_checkbox.set(True)
         self._ready_to_use_checkbox.set(True)
-        self._ready_to_use_button.click()
+        self._ready_to_use_button.click(timeout=10)
         self.wait_until_hidden()