Selenium Custom Field Decorator

AntepayFieldDecorator.java

package com.antepay.utility;

import org.openqa.selenium.WebElement;
import org.openqa.selenium.support.pagefactory.DefaultFieldDecorator;
import org.openqa.selenium.support.pagefactory.ElementLocator;
import org.openqa.selenium.support.pagefactory.ElementLocatorFactory;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Proxy;
import java.util.List;
import java.util.Map;

class AntepayFieldDecorator extends DefaultFieldDecorator {
    private Map<Class, ElementFactory> registeredElements;

    AntepayFieldDecorator(ElementLocatorFactory factory, Map<Class, ElementFactory> registeredElements) {
        super(factory);
        this.registeredElements = registeredElements;
    }

    @Override
    public Object decorate(ClassLoader loader, Field field) {
        if (!(WebElement.class.isAssignableFrom(field.getType())
            || isDecoratableList(field)
            || registeredElements.containsKey(field.getType()))) {
            return null;
        }

        ElementLocator locator = factory.createLocator(field);
        if (locator == null) {
            return null;
        }

        if (WebElement.class.isAssignableFrom(field.getType())) {
            return proxyForLocator(loader, locator);
        } else if (List.class.isAssignableFrom(field.getType())) {
            return proxyForListLocator(loader, locator);
        } else if (registeredElements.containsKey(field.getType())) {
            ElementFactory factory = registeredElements.get(field.getType());
            InvocationHandler handler = new LocatingElementHandler(locator, factory);
            return Proxy.newProxyInstance(loader, new Class[]{field.getType()}, handler);
        } else {
            return null;
        }
    }
}

LocatingElementHandler.java

package com.antepay.utility;

import org.openqa.selenium.NoSuchElementException;
import org.openqa.selenium.WrapsElement;
import org.openqa.selenium.support.pagefactory.ElementLocator;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

interface ElementFactory {
    WrapsElement create(ElementLocator locator);
}

public class LocatingElementHandler implements InvocationHandler {
    private final ElementLocator locator;
    private final ElementFactory elementFactory;

    LocatingElementHandler(ElementLocator locator, ElementFactory elementFactory) {
        this.locator = locator;
        this.elementFactory = elementFactory;
    }

    public Object invoke(Object object, Method method, Object[] objects) throws Throwable {
        WrapsElement element;
        try {
            element = elementFactory.create(locator);
        } catch (NoSuchElementException e) {
            if ("toString".equals(method.getName())) {
                return "Proxy element for: " + locator.toString();
            }
            throw e;
        }

        if ("getWrappedElement".equals(method.getName())) {
            return element;
        }

        try {
            return method.invoke(element, objects);
        } catch (InvocationTargetException e) {
            // Unwrap the underlying exception
            throw e.getCause();
        }
    }
}

Page.java

package com.antepay.utility;

import org.openqa.selenium.WebDriver;
import org.openqa.selenium.support.PageFactory;
import org.openqa.selenium.support.pagefactory.AjaxElementLocatorFactory;
import org.openqa.selenium.support.ui.ISelect;
import org.openqa.selenium.support.ui.Select;
import org.openqa.selenium.support.ui.WebDriverWait;

import java.util.Map;

public class Page {
    protected WebDriver driver;
    
    @FindBy(id = "prefix")
    private ISelect prefixElement = null;

    private static Map<Class, ElementFactory> elements = Map.of(
        ISelect.class, (locator) -> new Select(locator.findElement()),
        // here you can register more elements
        // YourElementInterface.class, lambda function that creates YourElement
    );

    protected Page(WebDriver driver) {
        this.driver = driver;
        final AjaxElementLocatorFactory factory = new AjaxElementLocatorFactory(this.driver, 10);
        final AntepayFieldDecorator decorator = new AntepayFieldDecorator(factory, elements);
        PageFactory.initElements(decorator, this);
    }
}