Capture WebDriver Screenshots When Running Parallel Tests With TestNG - selenium

Question

I am currently capturing screenshots on failure and success in TestNG by way of overriding the TestListenerAdapter methods onTestFailure, and onTestSuccess respectively. In order to do this you need to specify which driver you want to take a screenshot of.
My question: Is there a good way to capture screenshots when running tests in parallel on the method level?
In order to run tests in parallel, each individual test needs a unique driver instance. So, at any given time you have x number of driver instances running. When it comes time to capture a screenshot, how do you determine which driver to use?
Code excerpts below:
public class OnFailureListener extends TestListenerAdapter {
#Override
public void onTestFailure(ITestResult tr) {
Screenshots.captureScreenshot(tr);
super.onTestFailure(tr);
}
--
public static void captureScreenshot(ITestResult tr) {
WebDriver driver = TestClass.driver;
if (driver instanceof TakesScreenshot) {
String filename = "path/to/screenshot/file";
try {
File scrFile = ((TakesScreenshot)driver).getScreenshotAs(OutputType.FILE);
FileUtils.copyFile(scrFile, new File(filename));
} catch (IOException e) { e.printStackTrace(); }
}

Answer 1

If you create a base test class with access to the driver, then that driver will always be the correct driver
The following will achieve this;
All test classes must extend a simple base test class;
public asbtract baseTestCase() {
private WebDriver driver;
public WebDriver getDriver() {
return driver;
}
#BeforeMethod
public void createDriver() {
driver=XXXXDriver();
}
#AfterMethod
public void tearDownDriver() {
if (driver != null){
try{
driver.quit();
}
catch (WebDriverException e) {
System.out.println("***** CAUGHT EXCEPTION IN DRIVER TEARDOWN *****");
System.out.println(e);
}
}
}
In your listener, you need to access the base class:
public class ScreenshotListener extends TestListenerAdapter {
#Override
public void onTestFailure(ITestResult result){
Object currentClass = result.getInstance();
WebDriver webDriver = ((BaseTest) currentClass).getDriver();
if (webDriver != null){
File f = ((TakesScreenshot)webDriver).getScreenshotAs(OutputType.FILE);
//etc.
}
}
}