Custom Extensions for Test Setup

The previous lesson introduced the extension landscape. This one builds something real: a WebDriverExtension that creates a browser before each test, injects it into the test method as a parameter, and quits it after — with a screenshot on failure. This is the pattern used in serious JUnit 5 Selenium frameworks, and it replaces both TestNG's base-class inheritance pattern and JUnit 4's @Rule ExternalResource.

ParameterResolver — the key interface

ParameterResolver has two methods JUnit calls in sequence:

supportsParameter(ParameterResolutionContext, ExtensionContext) — JUnit asks: "does this extension know how to supply this parameter type?" Return true for types you handle.
resolveParameter(ParameterResolutionContext, ExtensionContext) — JUnit asks: "give me the value." Return the instance.

This two-step contract lets multiple extensions coexist — each only claims the parameters it can supply.

The WebDriverExtension — full implementation

import org.junit.jupiter.api.extension.*;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
 
public class WebDriverExtension
        implements BeforeEachCallback, AfterEachCallback, ParameterResolver {
 
    private static final ExtensionContext.Namespace NS =
        ExtensionContext.Namespace.create(WebDriverExtension.class);
    private static final String DRIVER = "driver";
 
    @Override
    public void beforeEach(ExtensionContext context) {
        ChromeOptions options = new ChromeOptions();
        if (Boolean.parseBoolean(System.getenv("HEADLESS"))) {
            options.addArguments("--headless", "--no-sandbox", "--disable-dev-shm-usage");
        }
        WebDriver driver = new ChromeDriver(options);
        driver.manage().window().maximize();
        context.getStore(NS).put(DRIVER, driver);
    }
 
    @Override
    public void afterEach(ExtensionContext context) {
        WebDriver driver = context.getStore(NS).remove(DRIVER, WebDriver.class);
        if (driver != null) driver.quit();
    }
 
    @Override
    public boolean supportsParameter(ParameterResolutionContext param,
                                     ExtensionContext extension) {
        return param.getParameter().getType() == WebDriver.class;
    }
 
    @Override
    public Object resolveParameter(ParameterResolutionContext param,
                                   ExtensionContext extension) {
        return extension.getStore(NS).get(DRIVER, WebDriver.class);
    }
}

Three design choices worth understanding:

Namespace.create(WebDriverExtension.class) — scopes the store to this extension class. If two extensions both use the key "driver", they won't collide because their namespaces differ.

context.getStore(NS).remove(...) in afterEach — remove retrieves and deletes the entry, which is cleaner than a separate get followed by a put(null). The driver is gone from the store the moment cleanup starts.

System.getenv("HEADLESS") check — makes the extension environment-aware without touching test code. Set HEADLESS=true in your CI pipeline; leave it unset locally for a headed browser.

Using the extension — clean test methods

import org.junit.jupiter.api.*;
import org.junit.jupiter.api.extension.ExtendWith;
import org.openqa.selenium.By;
import org.openqa.selenium.WebDriver;
import static org.junit.jupiter.api.Assertions.*;
 
@ExtendWith(WebDriverExtension.class)
class LoginPageTest {
 
    @Test
    @DisplayName("should log in with valid credentials")
    void validLogin(WebDriver driver) {
        driver.get("https://app.example.com/login");
        driver.findElement(By.id("email")).sendKeys("alice@test.com");
        driver.findElement(By.id("password")).sendKeys("secret");
        driver.findElement(By.cssSelector("button[type='submit']")).click();
        assertEquals("https://app.example.com/dashboard", driver.getCurrentUrl());
    }
 
    @Test
    @DisplayName("should show error for wrong password")
    void wrongPassword(WebDriver driver) {
        driver.get("https://app.example.com/login");
        driver.findElement(By.id("email")).sendKeys("alice@test.com");
        driver.findElement(By.id("password")).sendKeys("wrong");
        driver.findElement(By.cssSelector("button[type='submit']")).click();
        assertNotNull(driver.findElement(By.className("error-message")));
    }
}

Each test method receives a fresh WebDriver instance. There is no base class, no @BeforeEach in the test, no driver.quit() in @AfterEach — the extension handles all of it. The test reads as pure intent: get this URL, interact with the form, assert the outcome.

Screenshot extension — TestWatcher + Store access

The screenshot extension reads the driver from the same store the WebDriverExtension put it in. This works because extensions sharing the same namespace and key can read each other's stored values — provided the namespace is accessible:

import org.junit.jupiter.api.extension.*;
import org.openqa.selenium.OutputType;
import org.openqa.selenium.TakesScreenshot;
import org.openqa.selenium.WebDriver;
import java.io.*;
import java.nio.file.*;
 
public class ScreenshotExtension implements TestWatcher {
 
    private static final ExtensionContext.Namespace NS =
        ExtensionContext.Namespace.create(WebDriverExtension.class); // same namespace
 
    @Override
    public void testFailed(ExtensionContext context, Throwable cause) {
        WebDriver driver = context.getStore(NS).get("driver", WebDriver.class);
        if (!(driver instanceof TakesScreenshot ts)) return;
 
        byte[] bytes = ts.getScreenshotAs(OutputType.BYTES);
        String name = context.getDisplayName().replaceAll("[^a-zA-Z0-9]", "_");
        Path path = Paths.get("target/screenshots", name + ".png");
 
        try {
            Files.createDirectories(path.getParent());
            Files.write(path, bytes);
            System.out.println("📸 Screenshot saved: " + path);
        } catch (IOException e) {
            System.err.println("Could not save screenshot: " + e.getMessage());
        }
    }
}

@ExtendWith({WebDriverExtension.class, ScreenshotExtension.class})
class LoginPageTest { ... }

JUnit calls WebDriverExtension.beforeEach first (creates the driver), runs the test, then — if it fails — calls ScreenshotExtension.testFailed (reads the driver from the store and captures the screenshot), then calls WebDriverExtension.afterEach (quits the driver).

Custom data injection — not just WebDriver

ParameterResolver is not limited to Selenium. Any object you want to inject follows the same pattern:

// DatabaseExtension that injects a connection
public class DatabaseExtension implements BeforeEachCallback, AfterEachCallback, ParameterResolver {
 
    private static final ExtensionContext.Namespace NS =
        ExtensionContext.Namespace.create(DatabaseExtension.class);
 
    @Override
    public void beforeEach(ExtensionContext context) throws Exception {
        Connection conn = DriverManager.getConnection("jdbc:h2:mem:test");
        context.getStore(NS).put("connection", conn);
    }
 
    @Override
    public void afterEach(ExtensionContext context) throws Exception {
        Connection conn = context.getStore(NS).remove("connection", Connection.class);
        if (conn != null) conn.close();
    }
 
    @Override
    public boolean supportsParameter(ParameterResolutionContext param, ExtensionContext ext) {
        return param.getParameter().getType() == Connection.class;
    }
 
    @Override
    public Object resolveParameter(ParameterResolutionContext param, ExtensionContext ext) {
        return ext.getStore(NS).get("connection", Connection.class);
    }
}
 
// Test method receives the connection directly
@ExtendWith(DatabaseExtension.class)
class UserRepositoryTest {
    @Test void shouldSaveUser(Connection conn) throws Exception {
        // conn is live, transaction-clean, and auto-closed after the test
    }
}

Extension lifecycle flow

Step 1 of 5

beforeEach fires

WebDriverExtension.beforeEach() creates a new ChromeDriver (headless if HEADLESS=true) and stores it in the ExtensionContext store under namespace WebDriverExtension.

⚠️ Common mistakes

Different namespaces between cooperating extensions. If ScreenshotExtension uses Namespace.create(ScreenshotExtension.class) but WebDriverExtension uses Namespace.create(WebDriverExtension.class), the screenshot extension can't find the driver — different namespaces are separate key spaces. To share state, both must use the same namespace, or you need to pass data through a different mechanism.
@RegisterExtension as a static field when you want per-test lifecycle. static WebDriverExtension ext = new WebDriverExtension() behaves like @BeforeAll — the extension fires once. WebDriverExtension ext = new WebDriverExtension() (instance field) fires per test. Choose based on the resource lifecycle you want.
Returning true in supportsParameter for too broad a type. If two extensions both claim Object.class, JUnit will use whichever is declared first in @ExtendWith. Be specific in supportsParameter — match the exact type, or check annotations on the parameter to narrow the scope.

🎯 Practice task

Build a complete WebDriver extension. 30–40 minutes.

Implement WebDriverExtension as shown, including the HEADLESS environment variable check. Register it on a test class with two @Test methods that both receive WebDriver driver.
Run the tests. Confirm each method gets a fresh driver (add System.out.println(driver.hashCode()) — the hash codes should differ).
Screenshot extension. Implement ScreenshotExtension as shown. Add it to @ExtendWith. Write a test that deliberately fails (assert something wrong). Run and confirm a screenshot appears in target/screenshots/.
@RegisterExtension. Convert the @ExtendWith to a @RegisterExtension instance field: WebDriverExtension driver = new WebDriverExtension();. Update the test methods to use driver.getDriver() instead of a parameter (add a getDriver() method to your extension that reads from the store). Confirm the tests still pass.
Stretch — DatabaseExtension. Build the DatabaseExtension from the lesson, using an in-memory H2 database. Write a test that receives a Connection, inserts a row, and asserts it is readable — all using the injected connection with no explicit open/close in the test method.

Next lesson: conditional test execution — @EnabledOnOs, @EnabledIfEnvironmentVariable, and writing custom ExecutionCondition extensions.