Merge pull request #26 from browserbase/anirudh/e2e-selenium

Add examples in Selenium

Author: kamath

examples/__init__.py

@@ -8,9 +8,6 @@
 load_dotenv(override=True)
 
 # Make sure we have the required environment variables
-BROWSERBASE_CONNECT_URL = os.environ.get(
-    "BROWSERBASE_CONNECT_URL", "wss://connect.browserbase.com"
-)
 _BROWSERBASE_API_KEY = os.environ.get("BROWSERBASE_API_KEY")
 if not _BROWSERBASE_API_KEY:
     raise ValueError("BROWSERBASE_API_KEY is not set in environment")

examples/e2e/test_playwright.py

@@ -4,6 +4,7 @@
 import pytest
 from dotenv import load_dotenv
 from playwright.sync_api import Playwright, sync_playwright
+
 from browserbase import Browserbase
 
 from .. import (
@@ -20,9 +21,10 @@
 load_dotenv()
 
 CI = os.getenv("CI", "false").lower() == "true"
+MAX_RETRIES = 3
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(scope="function")  # Changed from "session" to "function"
 def playwright() -> Generator[Playwright, None, None]:
     with sync_playwright() as p:
         yield p

examples/e2e/test_selenium.py

@@ -0,0 +1,5 @@
+from .. import selenium_basic
+
+
+def test_selenium_basic() -> None:
+    selenium_basic.run()

examples/playwright_captcha.py

@@ -1,12 +1,6 @@
 from playwright.sync_api import Playwright, ConsoleMessage, sync_playwright
 
-
-from examples import (
-    BROWSERBASE_API_KEY,
-    BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
-    bb,
-)
+from examples import BROWSERBASE_PROJECT_ID, bb
 
 DEFAULT_CAPTCHA_URL = "https://www.google.com/recaptcha/api2/demo"
 OVERRIDE_TIMEOUT = 60000  # 60 seconds, adjust as needed
@@ -19,11 +13,7 @@ def run(playwright: Playwright) -> None:
     assert session.status == "RUNNING", f"Session status is {session.status}"
 
     # Connect to the remote session
-    connect_url = (
-        f"{BROWSERBASE_CONNECT_URL}?sessionId={session.id}&apiKey={BROWSERBASE_API_KEY}"
-    )
-    chromium = playwright.chromium
-    browser = chromium.connect_over_cdp(connect_url)
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
     context = browser.contexts[0]
     page = context.pages[0]
 

examples/playwright_contexts.py

@@ -1,21 +1,15 @@
 import time
 from typing import Optional
 
+from pydantic import TypeAdapter
 from playwright.sync_api import Cookie, Browser, Playwright, sync_playwright
 
-
-from examples import (
-    BROWSERBASE_API_KEY,
-    BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
-    bb,
-)
+from examples import BROWSERBASE_PROJECT_ID, bb
 from browserbase.types.session_create_params import (
     BrowserSettings,
     BrowserSettingsContext,
 )
 
-
 CONTEXT_TEST_URL = "https://www.browserbase.com"
 SECOND = 1000
 
@@ -47,10 +41,11 @@ def run(playwright: Playwright) -> None:
     # Step 2: Creates a session with the context
     session = bb.sessions.create(
         project_id=BROWSERBASE_PROJECT_ID,
-        browser_settings=BrowserSettings(
-            context=BrowserSettingsContext(id=context_id, persist=True),
+        browser_settings=TypeAdapter(BrowserSettings).validate_python(
+            {"context": {"id": context_id, "persist": True}}
         ),
     )
+    print(session)
 
     assert (
         session.context_id == context_id
@@ -59,9 +54,7 @@ def run(playwright: Playwright) -> None:
 
     # Step 3: Populates and persists the context
     print(f"Populating context {context_id} during session {session_id}")
-    connect_url = (
-        f"{BROWSERBASE_CONNECT_URL}?sessionId={session_id}&apiKey={BROWSERBASE_API_KEY}"
-    )
+    connect_url = session.connect_url
     browser = playwright.chromium.connect_over_cdp(connect_url)
     page = browser.contexts[0].pages[0]
 
@@ -108,9 +101,7 @@ def run(playwright: Playwright) -> None:
 
     # Step 5: Uses context to find previous state
     print(f"Reusing context {context_id} during session {session_id}")
-    connect_url = (
-        f"{BROWSERBASE_CONNECT_URL}?sessionId={session_id}&apiKey={BROWSERBASE_API_KEY}"
-    )
+    connect_url = session.connect_url
     browser = playwright.chromium.connect_over_cdp(connect_url)
     page = browser.contexts[0].pages[0]
 

examples/playwright_downloads.py

@@ -4,13 +4,7 @@
 
 from playwright.sync_api import Playwright, sync_playwright
 
-
-from examples import (
-    BROWSERBASE_API_KEY,
-    BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
-    bb,
-)
+from examples import BROWSERBASE_PROJECT_ID, bb
 
 download_re = re.compile(r"sandstorm-(\d{13})+\.mp3")
 
@@ -27,9 +21,7 @@ def run(playwright: Playwright) -> None:
     assert session.status == "RUNNING", f"Session status is {session.status}"
 
     # Connect to the remote session
-    connect_url = (
-        f"{BROWSERBASE_CONNECT_URL}?sessionId={session.id}&apiKey={BROWSERBASE_API_KEY}"
-    )
+    connect_url = session.connect_url
     browser = playwright.chromium.connect_over_cdp(connect_url)
     context = browser.contexts[0]
     page = context.pages[0]

examples/playwright_extensions.py

@@ -7,9 +7,7 @@
 from playwright.sync_api import Page, Playwright, sync_playwright
 
 from examples import (
-    BROWSERBASE_API_KEY,
     BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
     bb,
 )
 from browserbase.types import Extension, SessionCreateResponse
@@ -100,9 +98,7 @@ def run(playwright: Playwright) -> None:
         extension_id=extension.id,
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
     context = browser.contexts[0]
     page = context.pages[0]
     check_for_message(page, expected_message)
@@ -116,9 +112,7 @@ def run(playwright: Playwright) -> None:
         proxies=True,
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session_with_proxy.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session_with_proxy.connect_url)
     context = browser.contexts[0]
     page = context.pages[0]
 

examples/playwright_proxy.py

@@ -3,9 +3,7 @@
 from playwright.sync_api import Page, Playwright, sync_playwright
 
 from examples import (
-    BROWSERBASE_API_KEY,
     BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
     bb,
 )
 
@@ -26,9 +24,7 @@ def check_proxy_bytes(session_id: str) -> None:
 def run_enable_via_create_session(playwright: Playwright) -> None:
     session = bb.sessions.create(project_id=BROWSERBASE_PROJECT_ID, proxies=True)
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]
@@ -45,9 +41,7 @@ def run_enable_via_create_session(playwright: Playwright) -> None:
 def run_enable_via_querystring_with_created_session(playwright: Playwright) -> None:
     session = bb.sessions.create(project_id=BROWSERBASE_PROJECT_ID, proxies=True)
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}&enableProxy=true"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]
@@ -84,9 +78,7 @@ def run_geolocation_country(playwright: Playwright) -> None:
         ],
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]
@@ -113,9 +105,7 @@ def run_geolocation_state(playwright: Playwright) -> None:
         ],
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]
@@ -143,9 +133,7 @@ def run_geolocation_american_city(playwright: Playwright) -> None:
         ],
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]
@@ -172,9 +160,7 @@ def run_geolocation_non_american_city(playwright: Playwright) -> None:
         ],
     )
 
-    browser = playwright.chromium.connect_over_cdp(
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
 
     context = browser.contexts[0]
     page = context.pages[0]

examples/playwright_upload.py

@@ -2,12 +2,7 @@
 
 from playwright.sync_api import Playwright, sync_playwright
 
-from examples import (
-    BROWSERBASE_API_KEY,
-    BROWSERBASE_PROJECT_ID,
-    BROWSERBASE_CONNECT_URL,
-    bb,
-)
+from examples import BROWSERBASE_PROJECT_ID, bb
 
 PATH_TO_UPLOAD = Path.cwd() / "examples" / "packages" / "logo.png"
 
@@ -16,13 +11,8 @@ def run(playwright: Playwright) -> None:
     # Create a session
     session = bb.sessions.create(project_id=BROWSERBASE_PROJECT_ID)
 
-    # Construct the URL
-    url = (
-        f"{BROWSERBASE_CONNECT_URL}?apiKey={BROWSERBASE_API_KEY}&sessionId={session.id}"
-    )
-
     # Connect to the browser
-    browser = playwright.chromium.connect_over_cdp(url)
+    browser = playwright.chromium.connect_over_cdp(session.connect_url)
     context = browser.contexts[0]
     page = context.pages[0]
 

examples/selenium_basic.py

@@ -0,0 +1,59 @@
+from typing import Dict
+
+from selenium import webdriver
+from selenium.webdriver.remote.remote_connection import RemoteConnection
+
+from examples import BROWSERBASE_API_KEY, BROWSERBASE_PROJECT_ID, bb
+
+
+class BrowserbaseConnection(RemoteConnection):
+    """
+    Manage a single session with Browserbase.
+    """
+
+    session_id: str
+
+    def __init__(self, session_id: str, *args, **kwargs):  # type: ignore
+        super().__init__(*args, **kwargs)  # type: ignore
+        self.session_id = session_id
+
+    def get_remote_connection_headers(  # type: ignore
+        self, parsed_url: str, keep_alive: bool = False
+    ) -> Dict[str, str]:
+        headers = super().get_remote_connection_headers(parsed_url, keep_alive)  # type: ignore
+
+        # Update headers to include the Browserbase required information
+        headers["x-bb-api-key"] = BROWSERBASE_API_KEY
+        headers["session-id"] = self.session_id
+
+        return headers  # type: ignore
+
+
+def run() -> None:
+    # Use the custom class to create and connect to a new browser session
+    session = bb.sessions.create(project_id=BROWSERBASE_PROJECT_ID)
+    connection = BrowserbaseConnection(session.id, session.selenium_remote_url)
+    driver = webdriver.Remote(
+        command_executor=connection, options=webdriver.ChromeOptions()  # type: ignore
+    )
+
+    # Print a bit of info about the browser we've connected to
+    print(
+        "Connected to Browserbase",
+        f"{driver.name} version {driver.caps['browserVersion']}",  # type: ignore
+    )
+
+    try:
+        # Perform our browser commands
+        driver.get("https://www.sfmoma.org")
+        print(f"At URL: {driver.current_url} | Title: {driver.title}")
+        assert driver.current_url == "https://www.sfmoma.org/"
+        assert driver.title == "SFMOMA"
+
+    finally:
+        # Make sure to quit the driver so your session is ended!
+        driver.quit()
+
+
+if __name__ == "__main__":
+    run()