diff options
| -rw-r--r-- | python/report_ops/scotia.py | 27 |
1 files changed, 17 insertions, 10 deletions
diff --git a/python/report_ops/scotia.py b/python/report_ops/scotia.py index 0858af7e..c7297e7d 100644 --- a/python/report_ops/scotia.py +++ b/python/report_ops/scotia.py @@ -1,26 +1,33 @@ +import logging +import datetime +from contextlib import contextmanager import re from selenium import webdriver from selenium.webdriver.firefox.options import Options from selenium.webdriver.support.ui import WebDriverWait from selenium.webdriver.support import expected_conditions as EC from selenium.webdriver.common.by import By -import logging -import datetime from serenitas.utils.env import DAILY_DIR -from serenitas.analytics.dates import prev_business_day logger = logging.getLogger(__name__) +@contextmanager +def quit_driver(driver): + try: + yield driver + finally: + driver.quit() + + def download_report(account_username, password, report_dir, workdate): - driver = create_driver(report_dir) - login(driver, account_username, password) - navigate_to_inbox(driver, workdate) - attachment_element = get_attachment_element(driver) - attachment_element.click() - fname = attachment_element.text.split()[0] - driver.quit() + with quit_driver(create_driver(report_dir)) as driver: + login(driver, account_username, password) + navigate_to_inbox(driver, workdate) + attachment_element = get_attachment_element(driver) + attachment_element.click() + fname = attachment_element.text.split()[0] return fname |