[wwwsearch-commits] r28678 - in wwwsearch/mechanize/trunk: mechanize test

Sun Jun 11 19:25:23 CEST 2006

Author: jjlee
Date: Sun Jun 11 19:25:22 2006
New Revision: 28678

Modified:
   wwwsearch/mechanize/trunk/mechanize/_mechanize.py
   wwwsearch/mechanize/trunk/test/test_mechanize.py
Log:
Add a .set_cookie() method to Browser

Modified: wwwsearch/mechanize/trunk/mechanize/_mechanize.py
==============================================================================

--- wwwsearch/mechanize/trunk/mechanize/_mechanize.py	(original)
+++ wwwsearch/mechanize/trunk/mechanize/_mechanize.py	Sun Jun 11 19:25:22 2006
@@ -275,6 +275,20 @@
     def clear_history(self):
         self._history.clear()
 
+    def set_cookie(self, cookie_string):
+        if self._response is None:
+            raise BrowserStateError("not viewing any document")
+        if self.request.get_type() not in ["http", "https"]:
+            raise BrowserStateError("can't set cookie for non-HTTP/HTTPS "
+                                    "transactions")
+        cookiejar = self._ua_handlers["_cookies"].cookiejar
+        response = self.response()  # copy
+        headers = response.info()
+        if "Set-cookie" in headers:
+            del headers["Set-cookie"]
+        headers["Set-cookie"] = cookie_string
+        cookiejar.extract_cookies(response, self.request)
+
     def links(self, **kwds):
         """Return iterable over links (mechanize.Link objects)."""
         if not self.viewing_html():

Modified: wwwsearch/mechanize/trunk/test/test_mechanize.py
==============================================================================
--- wwwsearch/mechanize/trunk/test/test_mechanize.py	(original)
+++ wwwsearch/mechanize/trunk/test/test_mechanize.py	Sun Jun 11 19:25:22 2006
@@ -633,6 +633,51 @@
             r = b.open(url)
             self.assertEqual([link.absolute_url for link in b.links()], urls)
 
+    def test_set_cookie(self):
+        class CookieTestBrowser(TestBrowser):
+            default_features = list(TestBrowser.default_features)+["_cookies"]
+
+        # have to be visiting HTTP/HTTPS URL
+        url = "ftp://example.com/"
+        br = CookieTestBrowser()
+        r = mechanize.make_response(
+            "<html><head><title>Title</title></head><body></body></html>",
+            [("content-type", "text/html")],
+            url,
+            200, "OK",
+            )
+        br.add_handler(make_mock_handler()([("http_open", r)]))
+        handler = br._ua_handlers["_cookies"]
+        cj = handler.cookiejar
+        self.assertRaises(mechanize.BrowserStateError,
+                          br.set_cookie, "foo=bar")
+        self.assertEqual(len(cj), 0)
+
+
+        url = "http://example.com/"
+        br = CookieTestBrowser()
+        r = mechanize.make_response(
+            "<html><head><title>Title</title></head><body></body></html>",
+            [("content-type", "text/html")],
+            url,
+            200, "OK",
+            )
+        br.add_handler(make_mock_handler()([("http_open", r)]))
+        handler = br._ua_handlers["_cookies"]
+        cj = handler.cookiejar
+
+        # have to be visiting a URL
+        self.assertRaises(mechanize.BrowserStateError,
+                          br.set_cookie, "foo=bar")
+        self.assertEqual(len(cj), 0)
+
+
+        # normal case
+        br.open(url)
+        br.set_cookie("foo=bar")
+        self.assertEqual(len(cj), 1)
+        self.assertEqual(cj._cookies["example.com"]["/"]["foo"].value, "bar")
+
 
 class ResponseTests(TestCase):
     def test_set_response(self):
