From 368ae9b58dd77635fac012cce5518c3526ecb270 Mon Sep 17 00:00:00 2001
From: boatx <1446337+boatx@users.noreply.github.com>
Date: Thu, 22 Aug 2019 00:53:24 +0200
Subject: [PATCH] Convert duplicated getCookie method to helper function

---
 test_seasurf.py | 65 +++++++++++++------------------------------------
 1 file changed, 17 insertions(+), 48 deletions(-)

diff --git a/test_seasurf.py b/test_seasurf.py
index d81fe1f..b5e3f6f 100644
--- a/test_seasurf.py
+++ b/test_seasurf.py
@@ -25,6 +25,15 @@
     b = lambda s: s.encode('utf-8')
 
 
+def get_cookie(response, cookie_name):
+    cookies = response.headers.getlist('Set-Cookie')
+    for cookie in cookies:
+        value = parse_cookie(cookie).get(cookie_name)
+        if value:
+            return value
+    return None
+
+
 class BaseTestCase(unittest.TestCase):
     # Methods for backwards compatibility with python 2.5 & 2.6
     def assertIn(self, value, container, err=None):
@@ -271,7 +280,7 @@ def test_exempt_view(self):
         with self.app.test_client() as c:
             rv = c.post('/foo')
             self.assertIn(b('bar'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             self.assertEqual(cookie, None)
 
     def test_token_validation(self):
@@ -279,14 +288,6 @@ def test_token_validation(self):
         rv = self.app.test_client().post('/bar')
         self.assertIn(b('403 Forbidden'), rv.data)
 
-    def getCookie(self, response, cookie_name):
-        cookies = response.headers.getlist('Set-Cookie')
-        for cookie in cookies:
-            value = parse_cookie(cookie).get(cookie_name)
-            if value:
-                return value
-        return None
-
 
 class SeaSurfTestCaseIncludeViews(BaseTestCase):
     def setUp(self):
@@ -365,7 +366,7 @@ def test_exempt_view(self):
         with self.app.test_client() as c:
             rv = c.post('/foo/quz')
             self.assertIn(b('bar'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             self.assertEqual(cookie, None)
 
     def test_token_validation(self):
@@ -373,18 +374,10 @@ def test_token_validation(self):
             # should produce a logger warning
             rv = c.post('/bar')
             self.assertIn(b('403 Forbidden'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             token = self.csrf._get_token()
             self.assertEqual(cookie, token)
 
-    def getCookie(self, response, cookie_name):
-        cookies = response.headers.getlist('Set-Cookie')
-        for cookie in cookies:
-            value = parse_cookie(cookie).get(cookie_name)
-            if value:
-                return value
-        return None
-
 
 class SeaSurfTestCaseDisableCookie(unittest.TestCase):
     def setUp(self):
@@ -427,31 +420,23 @@ def test_has_csrf_cookie(self):
         with self.app.test_client() as c:
             rv = c.get('/foo/quz')
             self.assertIn(b('bar'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             token = self.csrf._get_token()
             self.assertEqual(cookie, token)
 
     def test_no_csrf_cookie(self):
         with self.app.test_client() as c:
             rv = c.get('/foo/baz')
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             self.assertEqual(cookie, None)
 
     def test_no_csrf_cookie_even_after_manually_validated(self):
         with self.app.test_client() as c:
             rv = c.post('/manual')
             self.assertIn(b('403 Forbidden'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             self.assertEqual(cookie, None)
 
-    def getCookie(self, response, cookie_name):
-        cookies = response.headers.getlist('Set-Cookie')
-        for cookie in cookies:
-            value = parse_cookie(cookie).get(cookie_name)
-            if value:
-                return value
-        return None
-
 
 class SeaSurfTestCaseSkipValidation(unittest.TestCase):
     def setUp(self):
@@ -497,7 +482,7 @@ def test_skips_validation(self):
         with self.app.test_client() as c:
             rv = c.post('/foo/quz')
             self.assertIn(b('bar'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             token = self.csrf._get_token()
             self.assertEqual(cookie, token)
 
@@ -519,14 +504,6 @@ def test_manual_validation(self):
             rv = c.post('/manual')
             self.assertIn(b('403 Forbidden'), rv.data)
 
-    def getCookie(self, response, cookie_name):
-        cookies = response.headers.getlist('Set-Cookie')
-        for cookie in cookies:
-            value = parse_cookie(cookie).get(cookie_name)
-            if value:
-                return value
-        return None
-
 
 class SeaSurfTestManualValidation(unittest.TestCase):
     def setUp(self):
@@ -553,18 +530,10 @@ def test_can_manually_validate_exempt_views(self):
         with self.app.test_client() as c:
             rv = c.post('/manual')
             self.assertIn(b('403 Forbidden'), rv.data)
-            cookie = self.getCookie(rv, self.csrf._csrf_name)
+            cookie = get_cookie(rv, self.csrf._csrf_name)
             token = self.csrf._get_token()
             self.assertEqual(cookie, token)
 
-    def getCookie(self, response, cookie_name):
-        cookies = response.headers.getlist('Set-Cookie')
-        for cookie in cookies:
-            value = parse_cookie(cookie).get(cookie_name)
-            if value:
-                return value
-        return None
-
 
 class SeaSurfTestCaseSave(BaseTestCase):
     def setUp(self):