allow to use set_cookie when HTTP_HOST is set

3 files changed, 14 insertions, 4 deletions

diff --git a/tests/test_app.py b/tests/test_app.py
index 3cb4fe7..ed442b7 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -199,6 +199,13 @@ class TestCookies(unittest.TestCase):
         app.get('/')
         app.reset()
 
+        app = webtest.TestApp(cookie_app,
+                              extra_environ={'HTTP_HOST': 'testserver'})
+        app.set_cookie('foo', 'bar')
+        app.set_cookie('fizz', ';bar=baz')  # Make sure we're escaping.
+        app.get('/')
+        app.reset()
+
     def test_preserves_cookies(self):
         def cookie_app(environ, start_response):
             req = Request(environ)
diff --git a/webtest/app.py b/webtest/app.py
index d1c94fb..736a865 100644
--- a/webtest/app.py
+++ b/webtest/app.py
@@ -69,13 +69,13 @@ class CookiePolicy(http_cookiejar.DefaultCookiePolicy):
     Domain=localhost."""
 
     def return_ok_domain(self, cookie, request):
-        if cookie.domain == '.localhost':
+        if cookie.domain.endswith(request.origin_req_host):
             return True
         return http_cookiejar.DefaultCookiePolicy.return_ok_domain(
             self, cookie, request)
 
     def set_ok_domain(self, cookie, request):
-        if cookie.domain == '.localhost':
+        if cookie.domain.endswith(request.origin_req_host):
             return True
         return http_cookiejar.DefaultCookiePolicy.set_ok_domain(
             self, cookie, request)
@@ -231,6 +231,9 @@ class TestApp(object):
         Sets a cookie to be passed through with requests.
 
         """
+        cookie_domain = self.extra_environ.get('HTTP_HOST', 'localhost')
+        cookie_domain = cookie_domain.split(':', 1)[0]
+        cookie_domain = '.' + cookie_domain
         value = escape_cookie_value(value)
         cookie = http_cookiejar.Cookie(
             version=0,
@@ -238,7 +241,7 @@ class TestApp(object):
             value=value,
             port=None,
             port_specified=False,
-            domain='.localhost',
+            domain=cookie_domain,
             domain_specified=True,
             domain_initial_dot=False,
             path='/',
diff --git a/webtest/utils.py b/webtest/utils.py
index 3eb7f90..adb01f3 100644
--- a/webtest/utils.py
+++ b/webtest/utils.py
@@ -97,7 +97,7 @@ class _RequestCookieAdapter(object):
     """
     def __init__(self, request):
         self._request = request
-        self.origin_req_host = request.host
+        self.origin_req_host = request.host.split(':')[0]
 
     def is_unverifiable(self):
         return True  # sure? Why not?