How to take a screenshot when a test is failed in JUnit5

Question

I am trying to write functionality that will take a screenshot of my browser in case a test fails. What I care about:

• take a screenshot and add it to the Allure2 report.
• tests must be run in parallel.
• information about Tear down in report

I tried to use JUnit API to implement Test Watcher. I have such a problem with it that when it passes a Driver to the testFailed method (just to do SS in the case of test failed) then, I can't use the @AfterEach annotation with closing the Driver because otherwise, the Driver will be NULL.

The BaseTest class from which the tests inherit contains annotations from JUnit5 and initializes the Driver here, and I would like to shut it down here. The class is also extended by CustomTestWatcher class which implements Test Watcher

@ExtendWith(CustomTestWatcher.class)
public class TestBase {

protected WebDriver driver;

@BeforeAll
public static void beforeAll() {
    AllureFeatures.generateEnvironmentVariablesForAllure(); // generate enviroments.xml for allure report purposes
}

@BeforeEach
public void setupDriver() {
    WebDriverThreadLocal.setDriver();
    this.driver = WebDriverThreadLocal.getDriver();
    driver.get(BASE_URL);
}

@AfterEach
public void closeDriver() {
    WebDriverThreadLocal.teardown();
}

CustomTestWatcher Class I don't like this solution. It duplicates the code in 4 places for each test status (disabled, successful, aborted, failed). I would prefer to use @AfterEach annotations, unfortunately, I have no idea how... Tear-down will also not be visible on the report.

public class CustomTestWatcher implements TestWatcher {
@Override
public void testDisabled(ExtensionContext context, Optional<String> reason) {
    WebDriverThreadLocal.teardown();
}
@Override
public void testSuccessful(ExtensionContext context) {
    WebDriverThreadLocal.teardown();
}
@Override
public void testAborted(ExtensionContext context, Throwable cause) {
    WebDriverThreadLocal.teardown();
}
@Override
public void testFailed(ExtensionContext context, Throwable cause) {
    ScreenShotCreator.takeScreenShot(WebDriverThreadLocal.getDriver()); // pass the driver and do SS
    WebDriverThreadLocal.teardown();
}

I also read that Allure has a TestLifecycleListener interface where you can override the BeforeTestStop method, looks like the perfect place to do SS, just again the question is how to implement it...