Merge pull request #316 from taizo/custom_headers

Added support for custom headers in PhantomJS driver.
cobrateam · May 9, 2014 · 34ead6f · 34ead6f
2 parents 57761ea + 9fa2b84
commit 34ead6f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 1 deletion.
diff --git a/splinter/driver/webdriver/phantomjs.py b/splinter/driver/webdriver/phantomjs.py
@@ -17,11 +17,15 @@ class WebDriver(BaseWebDriver):
     element_class = WebDriverElement
 
     def __init__(self, user_agent=None, load_images=True,
-                 desired_capabilities=None, wait_time=2, **kwargs):
+                 desired_capabilities=None, wait_time=2,
+                 custom_headers={}, **kwargs):
         capabilities = DesiredCapabilities.PHANTOMJS.copy()
         if user_agent is not None:
             capabilities['phantomjs.page.settings.userAgent'] = user_agent
         capabilities['phantomjs.page.settings.loadImages'] = load_images
+        if isinstance(custom_headers, dict):
+            for name, value in custom_headers.items():
+                capabilities['phantomjs.page.customHeaders.%s' % name] = value
         if desired_capabilities:
             capabilities.update(desired_capabilities)
 

diff --git a/tests/fake_webapp.py b/tests/fake_webapp.py
@@ -96,6 +96,11 @@ def upload_file():
         return '|'.join(buffer)
 
 
+@app.route('/headers', methods=['GET'])
+def request_headers():
+    return str(request.headers)
+
+
 @app.route('/foo')
 def foo():
     return "BAR!"

diff --git a/tests/test_webdriver_phantomjs.py b/tests/test_webdriver_phantomjs.py
@@ -53,3 +53,23 @@ def test_can_work_on_popups(self):
         # FIXME: Check https://github.com/detro/ghostdriver/issues/180 to see if
         # we can implement this test
         pass
+
+
+class PhantomJSBrowserTestWithCustomHeaders(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        custom_headers = {'X-Splinter-Customheaders-1': 'Hello',
+                          'X-Splinter-Customheaders-2': 'Bye'}
+        cls.browser = Browser("phantomjs", custom_headers=custom_headers)
+
+    def test_create_a_phantomjs_with_custom_headers(self):
+        self.browser.visit(EXAMPLE_APP + 'headers')
+        self.assertTrue(
+            self.browser.is_text_present('X-Splinter-Customheaders-1: Hello'))
+        self.assertTrue(
+            self.browser.is_text_present('X-Splinter-Customheaders-2: Bye'))
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.browser.quit()