Use Python @contextmanager decorator to start and stop WebDriver

One of the most frustrating simple annoyances with using Selenium is the need to manage the creation and destruction of your WebDriver instances.

If your configuration isn’t correct, it won’t start. And if your test does not complete successfully (or if you forget to close it down properly) you will have an orphaned browser window and a stray webdriver process, for example chromedriver or geckodriver

I can’t count the number of times I’ve had to go into the command line and type:

(on Windows)

taskkill /F /IM chromedriver.exe /T

(on Linux or Mac)

killall chromedriver

or

ps -ef | grep '[c]hromedriver' | awk '{print $2}' | xargs -l kill -9 

When starting a Remote WebDriver instance, you need the Selenium Server URL (or command executor) and Desired Capabilities

from selenium import webdriver

selenium_url = "http://localhost:4444"
capabilities = {"browserName": "chrome"}
driver = webdriver.Remote(command_executor=selenium_url, desired_capabilities=capabilities)

driver.quit()

And when you’re done, you need to make sure that you call

driver.quit()

The obvious solution is to make sure that you setup and teardown your webdriver instance. But in practice this means wrapping your code in try/except/finally blocks or some other mechanism

from selenium import webdriver
from selenium.common.exceptions import WebDriverException

try:
    driver = webdriver.Chrome()
except WebDriverException as e:
    print(e)
finally:
    driver.quit()

PyTest has a cool fixture mechanism, so you can do this:

from selenium import webdriver
import pytest

@pytest.fixture
def driver():
    print("starting webdriver")
    driver = webdriver.Remote(command_executor="http://localhost:4444", desired_capabilities={"browserName": "chrome"})
    
    yield driver # passes execution control to your test code

    print("stopping webdriver")
    driver.quit()

And your test can look as simple as this, by passing in the fixture as an argument to your test:

def test_with_webdriver(driver):
    driver.get("https://fijiaaron.wordpress.com");
    print(driver.title)

But what if you’re not using Pytest? What if you’re not actually testing, but using Selenium for process automation or data scraping?

Python has a couple of very useful tools that can help you manage your driver instance.

Using a decorator (much like the pytest fixture — which is, in fact, a decorator) — you can wrap a function and return another function. What Pytest is doing is creating a generator function — where you can use the yield statement to pass temporary control, similar to a return statement, but it works on a generator — which means you create a decorator that turns your simple setup and teardowm function into a generator that you can then yield to your automation.

It’s possible to do this yourself, but this is such a common pattern, that Python already has a library built in for doing this, contextlib: https://docs.python.org/3/library/contextlib.html

Specifically, we want to use the contextmanager decorator:

from contextlib import contextmanager

@contextmanager
def managed_resource(*args, **kwds):
    # Code to acquire resource, e.g.:
    resource = acquire_resource(*args, **kwds)
    try:
        yield resource
    finally:
        # Code to release resource, e.g.:
        release_resource(resource)

>>> with managed_resource(timeout=3600) as resource:
...     # Resource is released at the end of this block,
...     # even if code in the block raises an exception

You’ve probably seen something like this before when you read a file:

with open("file.txt") as file:
    for line in file.readlines():
        print(line)

This handles the opening and closing of the file handle.

You can do the same thing with a WebDriver instance:

from contextlib import contextmanager
from selenium import webdriver 

url = "https://fijiaaron.wordpress.com"

@contextmanager
def chromedriver(*args, **kwargs):
    print("starting webdriver")
    driver = webdriver.Chrome()    
    try: 
        yield driver
    finally:
        print("quitting webdriver")
        driver.quit()

with chromedriver() as driver:
    driver.get(url)
    print(driver.title)

You can even pass in arguments including your configuration capabilities to webdriver:

@contextmanager
def browser(*args, **kwargs):
    print(f"starting webdriver with {kwargs}")
    driver = webdriver.Remote(**kwargs)    
    try: 
        yield driver
    finally:
        print("quitting webdriver")
        driver.quit()

with browser(command_executor="http://localhost:4444", desired_capabilities={"browserName": "firefox"}) as driver:
    driver.get(url)
    print(driver.title)

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s