diff --git a/mechanicalsoup/stateful_browser.py b/mechanicalsoup/stateful_browser.py index 7329860..fc32ce7 100644 --- a/mechanicalsoup/stateful_browser.py +++ b/mechanicalsoup/stateful_browser.py @@ -8,6 +8,8 @@ from .form import Form from .utils import LinkNotFoundError +from requests.structures import CaseInsensitiveDict + class _BrowserState: def __init__(self, page=None, url=None, form=None, request=None): @@ -222,6 +224,18 @@ def select_form(self, selector="form", nr=0): return self.form + def _merge_referer(self, **kwargs): + """Helper function to set the Referer header in kwargs passed to + requests, if it has not already been overriden by the user.""" + + referer = self.url + headers = CaseInsensitiveDict(kwargs.get('headers', {})) + if referer is not None and 'Referer' not in headers: + kwargs = kwargs.copy() + headers['Referer'] = referer + kwargs['headers'] = headers + return kwargs + def submit_selected(self, btnName=None, update_state=True, *args, **kwargs): """Submit the form that was selected with :func:`select_form`. @@ -237,13 +251,7 @@ def submit_selected(self, btnName=None, update_state=True, """ self.form.choose_submit(btnName) - referer = self.url - if referer is not None: - if 'headers' in kwargs: - kwargs['headers']['Referer'] = referer - else: - kwargs['headers'] = {'Referer': referer} - + kwargs = self._merge_referer(**kwargs) resp = self.submit(self.__state.form, url=self.__state.url, *args, **kwargs) if update_state: diff --git a/tests/test_stateful_browser.py b/tests/test_stateful_browser.py index 3ea6f5f..44b93cd 100644 --- a/tests/test_stateful_browser.py +++ b/tests/test_stateful_browser.py @@ -446,6 +446,20 @@ def test_referer_submit(httpbin): assert actual_ref == ref +def test_referer_submit_override(httpbin): + browser = mechanicalsoup.StatefulBrowser() + ref = "https://example.com/my-referer" + ref_override = "https://example.com/override" + page = submit_form_headers.format(httpbin.url + "/headers") + browser.open_fake_page(page, url=ref) + browser.select_form() + response = browser.submit_selected(headers={"referer": ref_override}) + headers = response.json()["headers"] + referer = headers["Referer"] + actual_ref = re.sub('/*$', '', referer) + assert actual_ref == ref_override + + def test_referer_submit_headers(httpbin): browser = mechanicalsoup.StatefulBrowser() ref = "https://example.com/my-referer"