[java] only augment with matching augmenters

SeleniumHQ · Sep 17, 2021 · 9c9dccf · 9c9dccf
1 parent 0a5b2d8
commit 9c9dccf
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/java/src/org/openqa/selenium/remote/Augmenter.java b/java/src/org/openqa/selenium/remote/Augmenter.java
@@ -180,7 +180,7 @@ public WebDriver augment(WebDriver driver) {
       .annotateType(AnnotationDescription.Builder.ofType(Augmentable.class).build())
       .method(named("isAugmented")).intercept(FixedValue.value(true));
 
-    for (Augmentation<?> augmentation : augmentations) {
+    for (Augmentation<?> augmentation : matchingAugmenters) {
       Class<?> iface = augmentation.interfaceClass;
 
       Object instance = augmentation.implementation.apply(caps, execute);

diff --git a/java/test/org/openqa/selenium/remote/AugmenterTest.java b/java/test/org/openqa/selenium/remote/AugmenterTest.java
@@ -32,6 +32,7 @@
 import org.openqa.selenium.WebElement;
 import org.openqa.selenium.firefox.FirefoxOptions;
 import org.openqa.selenium.html5.LocationContext;
+import org.openqa.selenium.html5.WebStorage;
 import org.openqa.selenium.testing.UnitTests;
 
 import java.util.ArrayList;
@@ -100,6 +101,17 @@ public void shouldNotAddInterfaceWhenBooleanValueForItIsFalse() {
     assertThat(returned).isNotInstanceOf(HasMagicNumbers.class);
   }
 
+  @Test
+  public void shouldNotUseNonMatchingInterfaces() {
+    Capabilities caps = new ImmutableCapabilities("magic.numbers", true);
+    WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
+
+    WebDriver returned = getAugmenter()
+      .addDriverAugmentation("magic.numbers", HasMagicNumbers.class, (c, exe) -> () -> 42)
+      .augment(driver);
+    assertThat(returned).isNotInstanceOf(WebStorage.class);
+  }
+
   @Test
   public void shouldDelegateToHandlerIfAdded() {
     Capabilities caps = new ImmutableCapabilities("foo", true);