diff --git a/test/test.py b/test/test.py
index 5dbd908..d4fc1fc 100644
--- a/test/test.py
+++ b/test/test.py
@@ -16,6 +16,8 @@ class test_widget():
   def setup_method(self, method):
     options = Options()
     cwd = os.getcwd()
+    options.add_argument("--headless")
+    options.add_argument("--window-size=1280x1500")
     options.add_experimental_option("prefs", {"download.default_directory": cwd})
     self.driver = webdriver.Chrome(options=options)
     self.vars = {}
@@ -25,7 +27,6 @@ def teardown_method(self, method):
 
   def test_elements(self):
     self.driver.get("http://localhost:8383/voila/render/example.ipynb")
-    self.driver.set_window_size(1280, 1080)
     time.sleep(3)
 
     try:
@@ -36,10 +37,12 @@ def test_elements(self):
 
   def download_widget_image(self):
     self.driver.get("http://localhost:8383/voila/render/example.ipynb")
-    self.driver.set_window_size(1280, 720)
     time.sleep(3)
     self.driver.save_screenshot("widget-01.png")
-    self.driver.execute_script("window.scrollTo(0, 1000)")
+
+    #Move to the end of the page
+    self.driver.find_element(By.CLASS_NAME, 'bz-canvas3d').click()
+    self.driver.find_element(By.TAG_NAME, 'body').send_keys(Keys.END)
     time.sleep(3)
     self.driver.save_screenshot("widget-02.png")