мне нравится с селениумом работать как раз в ООП + контекстный менеджер.
создали драйвер, сохранили его в поле объекта.
методы - конкретные юзкейсы - залогиниться, вбить в форму то-то или проверить что-то ещё
в мейне тогда остаётся что-то такое
with MySelenium() as my_selenuim:
my_selenium.login(**login_data)
last_data = my_selenium.get_last_data()