Implement a set_cookie method in TestApp.

2 files changed, 121 insertions, 0 deletions

diff --git a/tests/test_app.py b/tests/test_app.py
index 474499f..db94c87 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -166,6 +166,27 @@ class TestPasteVariables(unittest.TestCase):
 
 class TestCookies(unittest.TestCase):
 
+    def test_set_cookie(self):
+        def cookie_app(environ, start_response):
+            req = Request(environ)
+            self.assertEqual(req.cookies['foo'], 'bar')
+            self.assertEqual(req.cookies['fizz'], ';bar=baz')
+
+            status = to_bytes("200 OK")
+            body = ''
+            headers = [
+                ('Content-Type', 'text/html'),
+                ('Content-Length', str(len(body))),
+            ]
+            start_response(status, headers)
+            return [to_bytes(body)]
+
+        app = webtest.TestApp(cookie_app)
+        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 b57814a..6a979e3 100644
--- a/webtest/app.py
+++ b/webtest/app.py
@@ -15,6 +15,7 @@ import json
 import random
 import fnmatch
 import mimetypes
+import cookielib
 
 from base64 import b64encode
 
@@ -85,6 +86,74 @@ class TestRequest(webob.BaseRequest):
     ResponseClass = TestResponse
 
 
+# A list of illegal characters in a cookie and the escaped equivalent.
+# Taken from Python's cookie module.
+COOKIE_ESCAPE_CHAR_MAP = {
+    '\000' : '\\000',  '\001' : '\\001',  '\002' : '\\002',
+    '\003' : '\\003',  '\004' : '\\004',  '\005' : '\\005',
+    '\006' : '\\006',  '\007' : '\\007',  '\010' : '\\010',
+    '\011' : '\\011',  '\012' : '\\012',  '\013' : '\\013',
+    '\014' : '\\014',  '\015' : '\\015',  '\016' : '\\016',
+    '\017' : '\\017',  '\020' : '\\020',  '\021' : '\\021',
+    '\022' : '\\022',  '\023' : '\\023',  '\024' : '\\024',
+    '\025' : '\\025',  '\026' : '\\026',  '\027' : '\\027',
+    '\030' : '\\030',  '\031' : '\\031',  '\032' : '\\032',
+    '\033' : '\\033',  '\034' : '\\034',  '\035' : '\\035',
+    '\036' : '\\036',  '\037' : '\\037',
+
+    # Because of the way browsers really handle cookies (as opposed
+    # to what the RFC says) we also encode , and ;
+
+    ',' : '\\054', ';' : '\\073',
+
+    '"' : '\\"',       '\\' : '\\\\',
+
+    '\177' : '\\177',  '\200' : '\\200',  '\201' : '\\201',
+    '\202' : '\\202',  '\203' : '\\203',  '\204' : '\\204',
+    '\205' : '\\205',  '\206' : '\\206',  '\207' : '\\207',
+    '\210' : '\\210',  '\211' : '\\211',  '\212' : '\\212',
+    '\213' : '\\213',  '\214' : '\\214',  '\215' : '\\215',
+    '\216' : '\\216',  '\217' : '\\217',  '\220' : '\\220',
+    '\221' : '\\221',  '\222' : '\\222',  '\223' : '\\223',
+    '\224' : '\\224',  '\225' : '\\225',  '\226' : '\\226',
+    '\227' : '\\227',  '\230' : '\\230',  '\231' : '\\231',
+    '\232' : '\\232',  '\233' : '\\233',  '\234' : '\\234',
+    '\235' : '\\235',  '\236' : '\\236',  '\237' : '\\237',
+    '\240' : '\\240',  '\241' : '\\241',  '\242' : '\\242',
+    '\243' : '\\243',  '\244' : '\\244',  '\245' : '\\245',
+    '\246' : '\\246',  '\247' : '\\247',  '\250' : '\\250',
+    '\251' : '\\251',  '\252' : '\\252',  '\253' : '\\253',
+    '\254' : '\\254',  '\255' : '\\255',  '\256' : '\\256',
+    '\257' : '\\257',  '\260' : '\\260',  '\261' : '\\261',
+    '\262' : '\\262',  '\263' : '\\263',  '\264' : '\\264',
+    '\265' : '\\265',  '\266' : '\\266',  '\267' : '\\267',
+    '\270' : '\\270',  '\271' : '\\271',  '\272' : '\\272',
+    '\273' : '\\273',  '\274' : '\\274',  '\275' : '\\275',
+    '\276' : '\\276',  '\277' : '\\277',  '\300' : '\\300',
+    '\301' : '\\301',  '\302' : '\\302',  '\303' : '\\303',
+    '\304' : '\\304',  '\305' : '\\305',  '\306' : '\\306',
+    '\307' : '\\307',  '\310' : '\\310',  '\311' : '\\311',
+    '\312' : '\\312',  '\313' : '\\313',  '\314' : '\\314',
+    '\315' : '\\315',  '\316' : '\\316',  '\317' : '\\317',
+    '\320' : '\\320',  '\321' : '\\321',  '\322' : '\\322',
+    '\323' : '\\323',  '\324' : '\\324',  '\325' : '\\325',
+    '\326' : '\\326',  '\327' : '\\327',  '\330' : '\\330',
+    '\331' : '\\331',  '\332' : '\\332',  '\333' : '\\333',
+    '\334' : '\\334',  '\335' : '\\335',  '\336' : '\\336',
+    '\337' : '\\337',  '\340' : '\\340',  '\341' : '\\341',
+    '\342' : '\\342',  '\343' : '\\343',  '\344' : '\\344',
+    '\345' : '\\345',  '\346' : '\\346',  '\347' : '\\347',
+    '\350' : '\\350',  '\351' : '\\351',  '\352' : '\\352',
+    '\353' : '\\353',  '\354' : '\\354',  '\355' : '\\355',
+    '\356' : '\\356',  '\357' : '\\357',  '\360' : '\\360',
+    '\361' : '\\361',  '\362' : '\\362',  '\363' : '\\363',
+    '\364' : '\\364',  '\365' : '\\365',  '\366' : '\\366',
+    '\367' : '\\367',  '\370' : '\\370',  '\371' : '\\371',
+    '\372' : '\\372',  '\373' : '\\373',  '\374' : '\\374',
+    '\375' : '\\375',  '\376' : '\\376',  '\377' : '\\377'
+    }
+
+
 class TestApp(object):
     """
     Wraps a WSGI application in a more convenient interface for
@@ -218,6 +287,37 @@ class TestApp(object):
     def cookies(self):
         return dict([(cookie.name, cookie.value) for cookie in self.cookiejar])
 
+    def set_cookie(self, name, value):
+        """
+        Sets a cookie to be passed through with requests.
+
+        """
+        # Quote the value
+        value = str(''.join(
+            COOKIE_ESCAPE_CHAR_MAP.get(x, x) for x in value
+        ))
+
+        # Set the cookie
+        cookie = cookielib.Cookie(
+            version=0,
+            name=name,
+            value=value,
+            port=None,
+            port_specified=False,
+            domain='.localhost',
+            domain_specified=True,
+            domain_initial_dot=False,
+            path='/',
+            path_specified=True,
+            secure=False,
+            expires=None,
+            discard=False,
+            comment=None,
+            comment_url=None,
+            rest=None
+        )
+        self.cookiejar.set_cookie(cookie)
+
     def reset(self):
         """
         Resets the state of the application; currently just clears