use virtual display

1 files changed, 24 insertions, 13 deletions

diff --git a/downloader.py b/downloader.py
index 58b5f42..7ab3c79 100644
--- a/downloader.py
+++ b/downloader.py
@@ -10,6 +10,7 @@ from selenium.common import TimeoutException
 from selenium.webdriver import ActionChains
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.wait import WebDriverWait
+from pyvirtualdisplay import Display
 
 from util import setupLogger
 
@@ -18,21 +19,31 @@ load_dotenv()
 logger = setupLogger()
 
 
-def driverSetup(url):
-    op = webdriver.ChromeOptions()
-    op.add_experimental_option('prefs', {
-        'download.default_directory': os.getenv('DOWNLOAD_PATH')
-    })
-    op.add_argument('--headless')
-    op.add_argument("--no-sandbox")
-    op.add_argument("--disable-dev-shm-usage")
-    driver = webdriver.Chrome(options=op)
-    driver.get(url)
-    return driver
+class Browser:
+    def __init__(self, url):
+        self.display = Display(visible=False, size=(800, 600))
+        self.display.start()
+        op = webdriver.ChromeOptions()
+        op.add_experimental_option('prefs', {
+            'download.default_directory': os.getenv('DOWNLOAD_PATH')
+        })
+        op.add_argument('--headless')
+        op.add_argument("--no-sandbox")
+        op.add_argument("--disable-dev-shm-usage")
+        self.driver = webdriver.Chrome(options=op)
+        self.driver.get(url)
+
+    def getDriver(self):
+        return self.driver
+
+    def close(self):
+        self.driver.quit()
+        self.display.stop()
 
 
 def downloadCsv():
-    driver = driverSetup(os.getenv('URL'))
+    browser = Browser(os.getenv('URL'))
+    driver = browser.getDriver()
     action = ActionChains(driver)
 
     try:
@@ -60,7 +71,7 @@ def downloadCsv():
         logger.error("Timed out while downloading report")
 
     finally:
-        driver.quit()
+        browser.close()
 
 
 if __name__ == '__main__':