Always wait for the page to load on PageObjects

So, I was just creating a simple selenium/JBehave code when a question appeared.

I'll first post the code simplified and then explain what is my question later.

So here we have a simple AbstractClass that will be inherited on my PageObjects. This class only contains a method to wait for certain elements on the page to be loaded. You can see how Im using it in the PageObject class (added a comment there).

AbstractPage.java

public abstract class AbstractPage {
    public void waitPageLoad() {
        WebDriverWait wait = new WebDriverWait(webDriverProvider.get(), 30);        
        wait.until(ExpectedConditions.visibilityOfAllElements(elementsToWait()));
    }

    protected List<WebElement> elementsToWait() {
        return null;
    }
}

PageObject.java

public class PageObject extends AbstractPage{
    @FindBy(id = "webElement1")
    private WebElement webElement1;
    
    @FindBy(id = "webElement2")
    private WebElement webElement2;

    public void clickWebElement1() {
        webElement1.click();
    }

    public void sendKeysWebElement2(String strKeys) {
        webElement2.sendKeys(strKeys);
    }

    //Note how im using the elementsToWait here
    @Override
    protected List<WebElement> elementsToWait() {
        return Arrays.asList(webElement1, webElement2);
    }
}

Now on my steps, if I want to wait for the page to load first and then do the actions that i want, I'd need to call the 'waitPageLoad()' method from my abstract class inside one of the steps (or all of them to be sure).

PageObjectSteps.java

@Component
public class PageObjectSteps {

    private PageObject pageObject;
    
    @When("User wants to click on webElement1")
    public void accountToDeposit () {
        pageObject.waitPageLoad(); //Calling here just as an example
        pageObject.clickWebElement1();
    }

    @When("User wants to type on webElement2 '$strType'")
    public void amountToDeposit(@Named("strType") String strType) {
        pageObject.sendKeysWebElement2(strType);
    }
}

Now my question is:

Is there a way that I could call the waitPageLoad() everytime that my pageObject is used but WITHOUT calling the method on the steps?

For example, I'd have a different waitPageLoad() for each pageObject, depending on what I'd need to wait for. On this example I'd wait for the webElement1 and webElement2 to be visible.

Does selenium have something like: @AlwaysWait where I could use before a method and it would be called everytime the page object is used(again, WITHOUT calling it in the steps)? Or a notation that would make a method to be called everytime the page object is used?

Example:

@AlwaysWait
public void waitPageObjectLoad() {
    WebDriverWait wait = new WebDriverWait(webDriverProvider.get(), 30);        
    wait.until(ExpectedConditions.visibilityOfAllElements(webElement1, webElement2));
}

Hopefully I made myself understandable, Thanks in advance.

PS: Asking around, I know that somehow you can use java reflections framework to do it, but I was wondering if you could do it with selenium only.

This is where you learn to love polymorphism, and the Proxy Pattern.

Create a new concrete class implementing the WebDriver interface called LazyWebDriver. Create two other classes to lazily load web elements: LazyWebElement and LazyWebElementList.

Methods in the LazyWebDriver should return LazyWebElement or LazyWebElementList objects, but the return values for those methods should be WebElement or List.

Now you just use the LazyWebDriver as if it were any other web driver. Finding elements using the standard WebDriver interface will always wait a certain number of seconds:

WebDriver driver = new ChromeDriver();
int secondsToWait = 15;
WebDriver lazyDriver = new LazyWebDriver(driver, secondsToWait);

// findElement(...) returns immediately
WebElement element = lazyDriver.findElement(By.id("foo"));

// Implicitly waits up to 15 seconds for the element
// to become visible before attempting to click on it
element.click();

// Returns immediately since the "wrapped" element
// has already been fetched after waiting.
String name = element.getAttribute("name");

The LazyWebDriver Class

public class LazyWebDriver implements WebDriver {
    private WebDriver driver;

    public LazyWebDriver(WebDriver driver, int secondsToWait) {
        this.driver = driver;
        this.wait = new WebDriverWait(driver, secondsToWait);
    }

    public void close() {
        driver.close();
    }

    public WebElement findElement(By by) {
        return new LazyWebElement(driver, by, wait);
    }

    public List<WebElement> findElements(By by) {
        return new LazyWebElementList(driver, by, wait);
    }

    // ... other methods just call through to driver.foo(...)
}

The LazyWebElement Class

public class LazyWebElement implements WebElement {
    private final WebDriver driver;
    private final WebDriverWait wait;
    private final By by;
    private WebElement element;

    public LazyWebElement(WebDriver driver, By by, WebDriverWait wait) {
        this.driver = driver;
        this.by = by;
        this.wait = wait;
    }

    private WebElement getElement() {
        if (element == null) {
            wait.until(ExpectedConditions.visibilityOfElementLocated(by));
            element = driver.findElement(by);
        }

        return element;
    }

    public void clear() {
        getElement().clear();
    }

    public void click() {
        getElement().click();
    }

    public String getAttribute(String attributeName) {
        return getElement().getAttribute(attributeName);
    }

    // Other methods in WebElement interface must first call getElement()
}

The LazyWebElementList Class

public class LazyWebElementList implements List<WebElement> {
    private final WebDriver driver;
    private final WebDriverWait wait;
    private final By by;
    private List<WebElement> elements;

    public LazyWebElementList(WebDriver driver, By by, WebDriverWait wait) {
        this.driver = driver;
        this.by = by;
        this.wait = wait;
    }

    private List<WebElement> getElements() {
        if (elements == null) {
            wait.until(ExpectedConditions.visibilityOfAllElementsLocated(by));
            elements = driver.findElements(by);
        }

        return elements;
    }

    public boolean add(WebElement element) {
        getElements().add(element);
    }

    public void clear() {
        getElements().clear();
    }

    // Other methods defined in List<E> interface must call getElements() first
}

I see in your code example you are getting the WebDriver object from a webDriverProvider. You can keep using this, except the web driver provider returns a LazyWebDriver cast to the WebDriver interface. The rest of your code remains completely ignorant that LazyWebDriver, LazyWebElement and LazyWebElementList even exist. This should be pretty easy to plug in to your existing test code.