[dotnet] Support assembly: Make IWebDriver extension methods handle w…

…rapped drivers
SeleniumHQ · Aug 26, 2021 · e571f89 · e571f89
1 parent 5f20319
commit e571f89
Showing 1 changed file with 22 additions and 2 deletions.
diff --git a/dotnet/src/support/Extensions/WebDriverExtensions.cs b/dotnet/src/support/Extensions/WebDriverExtensions.cs
@@ -37,7 +37,7 @@ public static class WebDriverExtensions
         /// indicate that it cannot take screenshots.</exception>
         public static Screenshot TakeScreenshot(this IWebDriver driver)
         {
-            ITakesScreenshot screenshotDriver = driver as ITakesScreenshot;
+            ITakesScreenshot screenshotDriver = GetDriverAs<ITakesScreenshot>(driver);
             if (screenshotDriver == null)
             {
                 IHasCapabilities capabilitiesDriver = driver as IHasCapabilities;
@@ -115,13 +115,33 @@ public static T ExecuteJavaScript<T>(this IWebDriver driver, string script, para
 
         private static object ExecuteJavaScriptInternal(IWebDriver driver, string script, object[] args)
         {
-            IJavaScriptExecutor executor = driver as IJavaScriptExecutor;
+            IJavaScriptExecutor executor = GetDriverAs<IJavaScriptExecutor>(driver);
             if (executor == null)
             {
                 throw new WebDriverException("Driver does not implement IJavaScriptExecutor");
             }
 
             return executor.ExecuteScript(script, args);
         }
+
+        private static T GetDriverAs<T>(IWebDriver driver) where T : class
+        {
+            T convertedDriver = driver as T;
+            if (convertedDriver == null)
+            {
+                // If the driver doesn't directly implement the desired interface, but does
+                // implement IWrapsDriver, walk up the hierarchy of wrapped drivers until
+                // either we find a class that does implement the desired interface, or is
+                // no longer wrapping a driver.
+                IWrapsDriver driverWrapper = driver as IWrapsDriver;
+                while (convertedDriver == null && driverWrapper != null)
+                {
+                    convertedDriver = driverWrapper.WrappedDriver as T;
+                    driverWrapper = driverWrapper.WrappedDriver as IWrapsDriver;
+                }
+            }
+
+            return convertedDriver;
+        }
     }
 }