1   package com.atlassian.pageobjects.elements.query.webdriver;
2   
3   import com.google.common.base.Function;
4   import com.google.common.base.Supplier;
5   import org.openqa.selenium.*;
6   
7   import static com.atlassian.pageobjects.elements.util.StringConcat.asString;
8   import static com.google.common.base.Preconditions.checkNotNull;
9   
10  /**
11   * <p>
12   * WebDriver based timed query that retrieves {@link org.openqa.selenium.WebElement} using provided
13   * {@link org.openqa.selenium.By} and applies provided function from that element to the target value.
14   *
15   * <p>
16   * If given element is not found, the 'invalid value' semantics of the timed query are applied. 
17   *
18   */
19  public class WebElementBasedTimedQuery<T> extends GenericWebDriverTimedQuery<T>
20  {
21      private final By by;
22  
23      public WebElementBasedTimedQuery(WebDriver driver, By by, final Function<WebElement, T> valueProvider, long timeout)
24      {
25          super(new WebElementSupplier<T>(driver, by, valueProvider), timeout);
26          this.by = by;
27      }
28  
29      public WebElementBasedTimedQuery(SearchContext searchContext, By by, final Function<WebElement, T> valueProvider, long timeout,
30              long interval)
31      {
32          super(new WebElementSupplier<T>(searchContext, by, valueProvider), timeout, interval);
33          this.by = by;
34      }
35  
36      public WebElementBasedTimedQuery(SearchContext searchContext, By by, final Function<WebElement, T> valueProvider, long timeout,
37              long interval, T invalidValue)
38      {
39          super(new WebElementSupplier<T>(searchContext, by, valueProvider, invalidValue), timeout, interval);
40          this.by = by;
41      }
42  
43      public WebElementBasedTimedQuery(WebElementBasedTimedQuery<T> origin, long timeout)
44      {
45          super(origin.webElementSupplier(), timeout, origin.interval);
46          this.by = null;
47      }
48  
49      WebElementSupplier<T> webElementSupplier()
50      {
51          return (WebElementSupplier<T>) valueSupplier; 
52      }
53  
54      @Override
55      public String toString()
56      {
57          return asString(super.toString(), "[locator=", by, "]");
58      }
59  
60      private static class WebElementSupplier<S> implements Supplier<S>
61      {
62          private final SearchContext searchContext;
63          private final By by;
64          private final Function<WebElement, S> valueProvider;
65          private final S invalidValue;
66  
67           public WebElementSupplier(final SearchContext searchContext, final By by, final Function<WebElement, S> valueProvider,
68                   S invalid)
69          {
70              this.valueProvider = valueProvider;
71              this.searchContext = checkNotNull(searchContext);
72              this.by = checkNotNull(by);
73              this.invalidValue = invalid;
74          }
75  
76          public WebElementSupplier(final SearchContext searchContext, final By by, final Function<WebElement, S> valueProvider)
77          {
78              this(searchContext, by, valueProvider, null);
79          }
80  
81          public S get()
82          {
83              try
84              {
85                  return valueProvider.apply(searchContext.findElement(by));
86              }
87              catch (NoSuchElementException e)
88              {
89                  throw new InvalidValue(invalidValue);
90              }
91          }
92      }
93  }