From c1d372003629c185f81125c5dab75fc3e8d13681 Mon Sep 17 00:00:00 2001
From: Lu Wang <coolwanglu@gmail.com>
Date: Tue, 17 Mar 2015 19:07:05 +0800
Subject: [PATCH] wait for page loaded in auto tests

---
 test/browser_tests.py       |  6 +++---
 test/test_local_browser.py  | 10 +++++++++-
 test/test_remote_browser.py | 13 +++++++++++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/test/browser_tests.py b/test/browser_tests.py
index e14b73a..278cbd8 100644
--- a/test/browser_tests.py
+++ b/test/browser_tests.py
@@ -28,7 +28,7 @@ class BrowserTests(Common):
     def tearDownClass(cls):
         pass
 
-    def run_test_case(self, filename, pdf2htmlEX_args=[]):
+    def run_test_case(self, filename, pdf2htmlEX_args=[], page_must_load=True):
         basefilename, extension = os.path.splitext(filename)
         htmlfilename = basefilename + '.html'
 
@@ -62,7 +62,7 @@ class BrowserTests(Common):
         out_img = Image.open(pngfilename_out_fullpath)
 
         pngfilename_ref_fullpath = os.path.join(png_out_dir, basefilename + '.ref.png')
-        self.generate_image(ref_htmlfilename, pngfilename_ref_fullpath)
+        self.generate_image(ref_htmlfilename, pngfilename_ref_fullpath, page_must_load=page_must_load)
         ref_img = Image.open(pngfilename_ref_fullpath)
 
         diff_img = ImageChops.difference(ref_img, out_img);
@@ -83,7 +83,7 @@ class BrowserTests(Common):
         # To test if the environment can detect any errors
         # E.g. when network is down, 404 message is shown for any HTML message
         with self.assertRaises(AssertionError):
-            self.run_test_case('test_fail.pdf')
+            self.run_test_case('test_fail.pdf', page_must_load=False)
 
     def test_basic_text(self):
         self.run_test_case('basic_text.pdf')
diff --git a/test/test_local_browser.py b/test/test_local_browser.py
index b317326..46527ff 100755
--- a/test/test_local_browser.py
+++ b/test/test_local_browser.py
@@ -5,6 +5,9 @@
 import unittest
 
 from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions
 from browser_tests import BrowserTests
 
 class test_local_browser(BrowserTests, unittest.TestCase):
@@ -24,6 +27,11 @@ class test_local_browser(BrowserTests, unittest.TestCase):
             cls.browser.quit()
         super(test_local_browser, cls).tearDownClass()
 
-    def generate_image(self, html_file, png_file):
+    def generate_image(self, html_file, png_file, page_must_load=True):
         self.browser.get('file://' + html_file)
+        try:
+            WebDriverWait(self.browser, 5).until(expected_conditions.presence_of_element_located((By.ID, 'page-container')))
+        except:
+            if page_must_load:
+                raise
         self.browser.save_screenshot(png_file)
diff --git a/test/test_remote_browser.py b/test/test_remote_browser.py
index 2857f4d..0bc7c40 100755
--- a/test/test_remote_browser.py
+++ b/test/test_remote_browser.py
@@ -7,6 +7,9 @@ import sys
 import os
 
 from selenium import webdriver
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support import expected_conditions
 from sauceclient import SauceClient
 from browser_tests import BrowserTests
 
@@ -95,7 +98,7 @@ class test_remote_browser_base(BrowserTests):
         try:
             passed = (sys.exc_info() == (None, None, None))
             branch = os.environ.get('TRAVIS_BRANCH', 'manual')
-            pull_request = os.environ.get('TRAVIS_PULL_REQUEST')
+            pull_request = os.environ.get('TRAVIS_PULL_REQUEST', 'false')
             self.sauce.jobs.update_job(self.browser.session_id, 
                 build_num=os.environ.get('TRAVIS_BUILD_NUMBER', '0'),
                 name='pdf2htmlEX',
@@ -104,10 +107,16 @@ class test_remote_browser_base(BrowserTests):
                 tags = [pull_request] if pull_request != 'false' else [branch]
             )
         except:
+            raise
             pass
 
-    def generate_image(self, html_file, png_file):
+    def generate_image(self, html_file, png_file, page_must_load=True):
         self.browser.get(BASEURL + html_file)
+        try:
+            WebDriverWait(self.browser, 5).until(expected_conditions.presence_of_element_located((By.ID, 'page-container')))
+        except:
+            if page_must_load:
+                raise
         self.browser.save_screenshot(png_file)
 
 test_classnames = []