test.py

from vfs_appointment_finder import utils

class TestUtils:
    
    def test_initialize_firefox_webdriver(self):
        fox = utils.initialize_webdriver('firefox')
        assert type(fox) == utils.webdriver.firefox.webdriver.WebDriver
        fox.quit()

    def test_initialize_chrome_webdriver(self):
        chrome = utils.initialize_webdriver('chrome')
        assert type(chrome) == utils.webdriver.chrome.webdriver.WebDriver
        chrome.quit()

utils.py

def initialize_webdriver(driver):
    """Initialize a webdriver and return it.
    :param browser: this specifies the type of the browser, either chrome or firefox.
                    If it is anything else than 'firefox' then a Chrome browser will be initalized.
    """
    if driver == 'firefox':
        options = webdriver.firefox.options.Options()
        options.headless = True
        return webdriver.Firefox(options=options)
    options = webdriver.chrome.options.Options()
    options.headless = True
    return webdriver.Chrome(options=options)