Cookie transfer fix

.github/workflows/codeql.yml

@@ -27,7 +27,7 @@ jobs:
           - python
 
     steps:
-      - uses: step-security/harden-runner@f4a75cfd619ee5ce8d5b864b0d183aff3c69b55a
+      - uses: step-security/harden-runner@95d9a5deda9de15063e7595e9719c11c38c90ae2
         with:
           disable-sudo: true
           egress-policy: block
@@ -36,15 +36,16 @@ jobs:
             github.com:443
             objects.githubusercontent.com:443
             uploads.github.com:443
+            release-assets.githubusercontent.com
 
-      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8
+      - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3
 
-      - uses: github/codeql-action/init@0499de31b99561a6d14a36a5f662c2a54f91beee
+      - uses: github/codeql-action/init@e12f0178983d466f2f6028f5cc7a6d786fd97f4b
         with:
           languages: ${{ matrix.language }}
 
-      - uses: github/codeql-action/autobuild@0499de31b99561a6d14a36a5f662c2a54f91beee
+      - uses: github/codeql-action/autobuild@e12f0178983d466f2f6028f5cc7a6d786fd97f4b
 
-      - uses: github/codeql-action/analyze@0499de31b99561a6d14a36a5f662c2a54f91beee
+      - uses: github/codeql-action/analyze@e12f0178983d466f2f6028f5cc7a6d786fd97f4b
         with:
           category: /language:${{matrix.language}}

requestium/requestium_mixin.py

@@ -7,7 +7,7 @@
 
 import tldextract
 from parsel.selector import Selector, SelectorList
-from selenium.common.exceptions import WebDriverException
+from selenium.common.exceptions import NoSuchWindowException, WebDriverException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.remote.webdriver import WebDriver as RemoteWebDriver
 from selenium.webdriver.support import expected_conditions
@@ -71,7 +71,7 @@ def ensure_add_cookie(self, cookie: dict[str, Any], override_domain: str | None
         cookie_domain = cookie["domain"] if cookie["domain"][0] != "." else cookie["domain"][1:]
         try:
             browser_domain = tldextract.extract(self.current_url).fqdn
-        except AttributeError:
+        except (AttributeError, NoSuchWindowException):
             browser_domain = ""
         if cookie_domain not in browser_domain:
             # TODO @joaqo: Check if hardcoding 'http' causes trouble.

tests/conftest.py

@@ -1,9 +1,11 @@
 # import os
 # import shutil
+import contextlib
 from typing import TYPE_CHECKING, cast
 
 import pytest
 from selenium import webdriver
+from selenium.common import WebDriverException
 
 import requestium
 
@@ -45,5 +47,9 @@ def session(request):  # noqa: ANN001, ANN201
         msg = f"Unknown driver type: {driver_type}"
         raise ValueError(msg)
 
-    with requestium.Session(driver=cast("DriverMixin", driver)) as session:
-        yield session
+    session = requestium.Session(driver=cast("DriverMixin", driver))
+    yield session
+
+    # Close all windows and end the session
+    with contextlib.suppress(WebDriverException, OSError):
+        driver.quit()

tests/test_cookies.py

@@ -4,12 +4,38 @@
 import requestium.requestium
 
 
-def test_transfer_session_cookies_to_driver(session: requestium.Session) -> None:
+@pytest.fixture(
+    params=[
+        {"name": "session_id", "value": "abc123", "domain": "example.com", "path": "/"},
+        {"name": "user_token", "value": "xyz789", "domain": "example.com", "path": "/"},
+    ],
+    ids=["session_id", "user_token"],
+)
+def cookie_data(request) -> dict[str, str]:  # noqa: ANN001
+    return request.param
+
+
+def test_transfer_driver_cookies_to_session(session: requestium.Session, cookie_data: dict[str, str]) -> None:
+    session.driver.get(f"https://{cookie_data['domain']}")
+    session.driver.add_cookie(cookie_data)
+
     assert session.cookies.keys() == []
-    response = session.get("http://google.com/")
-    assert response.cookies.keys().sort() == ["AEC", "NID"].sort()
+    session.transfer_driver_cookies_to_session()
+    assert session.cookies.keys() == [cookie_data["name"]]
+
+
+def test_transfer_session_cookies_to_driver(session: requestium.Session, cookie_data: dict[str, str]) -> None:
+    session.get(f"http://{cookie_data['domain']}")
+    session.cookies.set(name=cookie_data["name"], value=cookie_data["value"], domain=cookie_data["domain"], path=cookie_data["path"])
+
+    assert session.driver.get_cookies() == []
     session.transfer_session_cookies_to_driver()
-    assert session.cookies.keys() == ["AEC", "NID"]
+    driver_cookies = session.driver.get_cookies()
+    assert len(driver_cookies) == 1
+    assert driver_cookies[0]["name"] == cookie_data["name"]
+    assert driver_cookies[0]["value"] == cookie_data["value"]
+    assert driver_cookies[0]["domain"] in [cookie_data["domain"], f".{cookie_data['domain']}"]
+    assert driver_cookies[0]["path"] == cookie_data["path"]
 
 
 def test_transfer_session_cookies_to_driver_no_domain_error(session: requestium.Session) -> None: