[java] Fixing Augmenter to add default set of augmentations

quehanwei · Jul 17, 2020 · e1f5d29 · e1f5d29
1 parent b23fc08
commit e1f5d29
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 15 deletions.
diff --git a/java/client/src/org/openqa/selenium/remote/Augmenter.java b/java/client/src/org/openqa/selenium/remote/Augmenter.java
@@ -46,6 +46,7 @@
 import java.util.function.BiFunction;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
 import static net.bytebuddy.matcher.ElementMatchers.anyOf;
@@ -63,15 +64,24 @@ public class Augmenter {
   private final Set<Augmentation<?>> augmentations;
 
   public Augmenter() {
-    this.augmentations = Collections.emptySet();
-    addDriverAugmentation(new AddApplicationCache());
-    addDriverAugmentation(new AddLocationContext());
-    addDriverAugmentation(new AddNetworkConnection());
-    addDriverAugmentation(new AddRotatable());
-    addDriverAugmentation(new AddWebStorage());
+    augmentations = new HashSet<>();
+    Stream.of(
+        new AddApplicationCache(),
+        new AddLocationContext(),
+        new AddNetworkConnection(),
+        new AddRotatable(),
+        new AddWebStorage()
+    ).forEach(provider -> augmentations.add(createAugmentation(provider)));
 
     StreamSupport.stream(ServiceLoader.load(AugmenterProvider.class).spliterator(), false)
-      .forEach(this::addDriverAugmentation);
+        .forEach(provider -> augmentations.add(createAugmentation(provider)));
+  }
+
+  private static <X> Augmentation<X> createAugmentation(AugmenterProvider<X> provider) {
+    Require.nonNull("Interface provider", provider);
+    return new Augmentation<>(provider.isApplicable(),
+                              provider.getDescribedInterface(),
+                              provider::getImplementation);
   }
 
   private Augmenter(Set<Augmentation<?>> augmentations, Augmentation<?> toAdd) {

diff --git a/java/client/test/org/openqa/selenium/remote/AugmenterTest.java b/java/client/test/org/openqa/selenium/remote/AugmenterTest.java
@@ -30,6 +30,7 @@
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
 import org.openqa.selenium.firefox.FirefoxOptions;
+import org.openqa.selenium.html5.LocationContext;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -44,10 +45,32 @@
 
 public class AugmenterTest {
 
-  protected Augmenter getAugmenter() {
+  private Augmenter getAugmenter() {
     return new Augmenter();
   }
 
+  @Test
+  public void shouldAugmentRotatable() {
+    final Capabilities caps = new ImmutableCapabilities(CapabilityType.ROTATABLE, true);
+    WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
+
+    WebDriver returned = getAugmenter().augment(driver);
+
+    assertThat(returned).isNotSameAs(driver);
+    assertThat(returned).isInstanceOf(Rotatable.class);
+  }
+
+  @Test
+  public void shouldAugmentLocationContext() {
+    final Capabilities caps = new ImmutableCapabilities(CapabilityType.SUPPORTS_LOCATION_CONTEXT, true);
+    WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
+
+    WebDriver returned = getAugmenter().augment(driver);
+
+    assertThat(returned).isNotSameAs(driver);
+    assertThat(returned).isInstanceOf(LocationContext.class);
+  }
+
   @Test
   public void shouldAddInterfaceFromCapabilityIfNecessary() {
     final Capabilities caps = new ImmutableCapabilities("magic.numbers", true);
@@ -182,20 +205,15 @@ public interface FindByMagic {
 
   @Test
   public void shouldBeAbleToAugmentMultipleTimes() {
-    Capabilities caps = new ImmutableCapabilities("canRotate", true, "magic.numbers", true);
+    Capabilities caps = new ImmutableCapabilities("rotatable", true, "magic.numbers", true);
 
     StubExecutor stubExecutor = new StubExecutor(caps);
     stubExecutor.expect(DriverCommand.GET_SCREEN_ORIENTATION,
       Collections.emptyMap(),
       ScreenOrientation.PORTRAIT.name());
     RemoteWebDriver driver = new RemoteWebDriver(stubExecutor, caps);
 
-    WebDriver augmented = getAugmenter()
-      .addDriverAugmentation(
-        "canRotate",
-        Rotatable.class,
-        (c, exe) -> new RemoteRotatable(exe))
-      .augment(driver);
+    WebDriver augmented = getAugmenter().augment(driver);
 
     assertThat(driver).isNotSameAs(augmented);
     assertThat(augmented).isInstanceOf(Rotatable.class);