diff --git a/tests/viz/altair/test_saving.py b/tests/viz/altair/test_saving.py index 538cc9b..426ffae 100644 --- a/tests/viz/altair/test_saving.py +++ b/tests/viz/altair/test_saving.py @@ -1,6 +1,7 @@ import shutil from pathlib import Path from nesta_ds_utils.viz.altair import saving +from selenium.webdriver.chromium.webdriver import ChromiumDriver import pandas as pd import altair as alt import pytest @@ -41,3 +42,18 @@ def test_save_altair_exception(): saving.save( fig, "test_fig", path, save_png=False, save_html=False, save_svg=False ) + + +def test_webdriver(): + """Test that Chrome WebDriver is created by default is a ChromiumDriver, and context manager stops the webdriver.""" + driver = saving._google_chrome_driver_setup() + assert isinstance(driver, ChromiumDriver) + + with saving.webdriver_context(driver) as some_driver: + # No actions needed here, + # just testing that the context manager calls .quit() on driver to terminate. + pass + + # If subprocess not terminated, .poll() returns None + # https://docs.python.org/3/library/subprocess.html#subprocess.Popen.returncode + assert driver.service.process.poll() is not None