From 7b338ec6a895130d9594edc724461cd74bc12124 Mon Sep 17 00:00:00 2001
From: Sri Harsha <Harsha509@users.noreply.github.com>
Date: Tue, 17 Mar 2020 20:47:07 +0530
Subject: [PATCH] Add: Same site cookie feature in python bindings (#8114)

* Add: Adding sameSite cookie attribute to python bindings

Co-authored-by: David Burns <david.burns@theautomatedtester.co.uk>
---
 py/selenium/webdriver/remote/webdriver.py     |  9 +++--
 .../selenium/webdriver/common/cookie_tests.py | 36 ++++++++++++++++++-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/py/selenium/webdriver/remote/webdriver.py b/py/selenium/webdriver/remote/webdriver.py
index 73f09b1cdeb76..4ec2c680072c9 100644
--- a/py/selenium/webdriver/remote/webdriver.py
+++ b/py/selenium/webdriver/remote/webdriver.py
@@ -947,15 +947,20 @@ def add_cookie(self, cookie_dict):
 
         :Args:
          - cookie_dict: A dictionary object, with required keys - "name" and "value";
-            optional keys - "path", "domain", "secure", "expiry"
+            optional keys - "path", "domain", "secure", "expiry", "sameSite"
 
         Usage:
             driver.add_cookie({'name' : 'foo', 'value' : 'bar'})
             driver.add_cookie({'name' : 'foo', 'value' : 'bar', 'path' : '/'})
             driver.add_cookie({'name' : 'foo', 'value' : 'bar', 'path' : '/', 'secure':True})
+            driver.add_cookie({'name': 'foo', 'value': 'bar', 'sameSite': 'Strict'})
 
         """
-        self.execute(Command.ADD_COOKIE, {'cookie': cookie_dict})
+        if 'sameSite' in cookie_dict:
+          assert cookie_dict['sameSite'] in ['Strict', 'Lax']
+          self.execute(Command.ADD_COOKIE, {'cookie': cookie_dict})
+        else:
+          self.execute(Command.ADD_COOKIE, {'cookie': cookie_dict})
 
     # Timeouts
     def implicitly_wait(self, time_to_wait):
diff --git a/py/test/selenium/webdriver/common/cookie_tests.py b/py/test/selenium/webdriver/common/cookie_tests.py
index a1c0970fe7aba..412c179a55840 100644
--- a/py/test/selenium/webdriver/common/cookie_tests.py
+++ b/py/test/selenium/webdriver/common/cookie_tests.py
@@ -21,6 +21,7 @@
 
 import pytest
 
+from selenium.common.exceptions import WebDriverException
 
 @pytest.fixture
 def cookie(webserver):
@@ -32,6 +33,27 @@ def cookie(webserver):
         'secure': False}
     return cookie
 
+@pytest.fixture
+def same_site_cookie_strict(webserver):
+    same_site_cookie_strict = {
+        'name': 'foo',
+        'value': 'bar',
+        'path': '/',
+        'domain': webserver.host,
+        'sameSite': 'Strict',
+        'secure': False}
+    return same_site_cookie_strict
+
+@pytest.fixture
+def same_site_cookie_lax(webserver):
+    same_site_cookie_lax = {
+            'name': 'foo',
+            'value': 'bar',
+            'path': '/',
+            'domain': webserver.host,
+            'sameSite': 'Lax',
+            'secure': False}
+    return same_site_cookie_lax
 
 @pytest.fixture(autouse=True)
 def pages(request, driver, pages):
@@ -39,12 +61,24 @@ def pages(request, driver, pages):
     yield pages
     driver.delete_all_cookies()
 
-
 def testAddCookie(cookie, driver):
     driver.add_cookie(cookie)
     returned = driver.execute_script('return document.cookie')
     assert cookie['name'] in returned
 
+@pytest.mark.xfail_firefox(raises=WebDriverException,
+                           reason='sameSite cookie attribute not implemented')
+def testAddCookieSameSiteStrict(same_site_cookie_strict, driver):
+    driver.add_cookie(same_site_cookie_strict)
+    returned = driver.get_cookie('foo')
+    assert returned['sameSite'] == 'Strict'
+
+@pytest.mark.xfail_firefox(raises=WebDriverException,
+                           reason='sameSite cookie attribute not implemented')
+def testAddCookieSameSiteLax(same_site_cookie_lax, driver):
+    driver.add_cookie(same_site_cookie_lax)
+    returned = driver.get_cookie('foo')
+    assert returned['sameSite'] == 'Lax'
 
 @pytest.mark.xfail_ie
 def testAddingACookieThatExpiredInThePast(cookie, driver):