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