diff --git a/src/test/java/tests/BaseTest.java b/src/test/java/tests/BaseTest.java
index f560e15..bc7037a 100644
--- a/src/test/java/tests/BaseTest.java
+++ b/src/test/java/tests/BaseTest.java
@@ -4,6 +4,8 @@
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.firefox.FirefoxDriver;
 import org.openqa.selenium.firefox.FirefoxOptions;
+import org.testng.ITestContext;
+import org.testng.ITestResult;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
@@ -12,7 +14,7 @@
 import utils.logs.Log;
 
 public class BaseTest {
-    public WebDriver driver;
+    private WebDriver driver;
     public HomePage homePage;
 
     public WebDriver getDriver() {
@@ -20,12 +22,12 @@ public WebDriver getDriver() {
     }
 
     @BeforeSuite
-    public void setupProxy() {
+    public void setupProxy(ITestContext iTestContext) {
 
     }
 
     @BeforeClass
-    public void classLevelSetup() {
+    public void classLevelSetup(ITestContext iTestContext) {
         Log.info("Tests is starting!");
 
         FirefoxOptions firefoxOptions = new FirefoxOptions();
@@ -36,6 +38,7 @@ public void classLevelSetup() {
 
         firefoxOptions.setCapability("proxy", proxy);
         driver = new FirefoxDriver(firefoxOptions);
+        iTestContext.setAttribute("WebDriver", driver);
     }
 
     @BeforeMethod
diff --git a/src/test/java/utils/listeners/TestListener.java b/src/test/java/utils/listeners/TestListener.java
index c98e2b5..226a5ca 100644
--- a/src/test/java/utils/listeners/TestListener.java
+++ b/src/test/java/utils/listeners/TestListener.java
@@ -7,7 +7,6 @@
 import org.testng.ITestContext;
 import org.testng.ITestListener;
 import org.testng.ITestResult;
-import tests.BaseTest;
 import utils.extentreports.ExtentManager;
 import utils.logs.Log;
 
@@ -15,7 +14,7 @@
 
 import static utils.extentreports.ExtentTestManager.getTest;
 
-public class TestListener extends BaseTest implements ITestListener {
+public class TestListener implements ITestListener {
     private static String getTestMethodName(ITestResult iTestResult) {
         return iTestResult.getMethod().getConstructorOrMethod().getName();
     }
@@ -23,7 +22,6 @@ private static String getTestMethodName(ITestResult iTestResult) {
     @Override
     public void onStart(ITestContext iTestContext) {
         Log.info("I am in onStart method " + iTestContext.getName());
-        iTestContext.setAttribute("WebDriver", this.driver);
     }
 
     @Override
@@ -49,17 +47,28 @@ public void onTestSuccess(ITestResult iTestResult) {
     public void onTestFailure(ITestResult iTestResult) {
         Log.info(getTestMethodName(iTestResult) + " test is failed.");
 
-        //Get driver from BaseTest and assign to local webdriver variable.
-        Object testClass = iTestResult.getInstance();
-        WebDriver driver = ((BaseTest) testClass).getDriver();
+        //Get an instance of WebDriver out of the ITestResult
+        if (iTestResult.getTestContext().getAttribute("WebDriver") != null) {
+            WebDriver driver =
+                    (WebDriver) iTestResult.getTestContext()
+                            .getAttribute("WebDriver");
 
-        //Take base64Screenshot screenshot for extent reports
-        String base64Screenshot =
-                "data:image/png;base64," + ((TakesScreenshot) Objects.requireNonNull(driver)).getScreenshotAs(OutputType.BASE64);
+            //Take base64Screenshot screenshot for extent reports
+            String base64Screenshot =
+                    "data:image/png;base64," +
+                            ((TakesScreenshot) Objects.requireNonNull(driver))
+                            .getScreenshotAs(OutputType.BASE64);
 
-        //ExtentReports log and screenshot operations for failed tests.
-        getTest().log(Status.FAIL, "Test Failed",
-                getTest().addScreenCaptureFromBase64String(base64Screenshot).getModel().getMedia().get(0));
+            //ExtentReports log and screenshot operations for failed tests.
+            getTest().log(Status.FAIL, "Test Failed",
+                    getTest().addScreenCaptureFromBase64String(base64Screenshot)
+                            .getModel().getMedia().get(0));
+        } else {
+            throw new RuntimeException(
+                    "Unable to get a WebDriver instance out of the ITestResult. " +
+                    " Your Test class is supposed to set it by calling " +
+                            "ITestContext.addAttribute(\"WebDriver\", driver)");
+        }
     }
 
     @Override