SurfSense/surfsense_backend/scripts/e2e_google_search.py

"""Live end-to-end checks for the Google Search scraper (needs proxy + browser).

    .venv/Scripts/python.exe scripts/e2e_google_search.py

Covers: a plain query, a site: filter, text ads, product ads, the
focusOnPaidAds retry (commercial = ads found; non-commercial = retries capped,
organic still returned), People-Also-Ask answer expansion, sitelinks, the AI
Overview, the mobile layout, filter=0, base64 icons, and Google AI Mode.

Pass case names as args to run a subset, e.g.:

    .venv/Scripts/python.exe scripts/e2e_google_search.py paa
"""

import asyncio
import logging
import sys
import time
from pathlib import Path

from dotenv import load_dotenv

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_ROOT))
load_dotenv(_ROOT / ".env")

logging.basicConfig(level=logging.WARNING)
logging.getLogger("app.proprietary.platforms.google_search.scraper").setLevel(
    logging.INFO
)

from app.proprietary.platforms.google_search import (  # noqa: E402
    GoogleSearchScrapeInput,
    scrape_serps,
)
from app.proprietary.platforms.google_search.fetch import close_sessions  # noqa: E402


async def run_ai_mode(label: str, *, queries: str) -> None:
    print(f"\n=== {label} ===")
    t0 = time.perf_counter()
    inp = GoogleSearchScrapeInput(
        queries=queries,
        countryCode="us",
        languageCode="en",
        aiModeSearch={"enableAiMode": True},
    )
    items = await scrape_serps(inp, limit=2)
    ai_items = [i for i in items if i["aiModeResult"]]
    assert ai_items, f"{label}: no aiModeResult item emitted"
    res = ai_items[0]["aiModeResult"]
    print(
        f"  text={len(res['text'])} chars, sources={len(res['sources'])} "
        f"({time.perf_counter() - t0:.0f}s)"
    )
    print(f"  {res['text'][:130]!r}")
    for s in res["sources"][:3]:
        print(f"    src: {(s['title'] or '')[:60]!r}")
    assert res["text"] and len(res["text"]) > 100, f"{label}: answer too short"
    assert res["sources"], f"{label}: no cited sources"
    assert "udm=50" in ai_items[0]["searchQuery"]["url"]


async def run(
    label: str,
    *,
    expect_ads=False,
    expect_products=False,
    expect_paa_answers=False,
    expect_sitelinks=False,
    expect_aio=False,
    expect_device=None,
    expect_icons=False,
    **kwargs,
) -> None:
    print(f"\n=== {label} ===")
    t0 = time.perf_counter()
    inp = GoogleSearchScrapeInput(countryCode="us", languageCode="en", **kwargs)
    items = await scrape_serps(inp, limit=1)
    assert items, f"{label}: no SERP item"
    it = items[0]
    paa_answered = [p for p in it["peopleAlsoAsk"] if p["answer"]]
    sitelinked = [o for o in it["organicResults"] if o["siteLinks"]]
    print(f"  term={it['searchQuery']['term']!r} resultsTotal={it['resultsTotal']}")
    print(
        f"  organic={len(it['organicResults'])} paidResults={len(it['paidResults'])} "
        f"paidProducts={len(it['paidProducts'])} related={len(it['relatedQueries'])} "
        f"suggested={len(it['suggestedResults'])} "
        f"paa={len(it['peopleAlsoAsk'])} (answered={len(paa_answered)}) "
        f"({time.perf_counter() - t0:.0f}s)"
    )
    for o in sitelinked[:2]:
        print(
            f"    [sitelinks on #{o['position']}] "
            + ", ".join(s["title"] for s in o["siteLinks"][:5])
        )
    aio = it["aiOverview"]
    if aio:
        print(
            f"    [aiOverview] content={len(aio['content'])} chars, "
            f"sources={len(aio['sources'])}"
        )
        print(f"        {aio['content'][:110]!r}")
        for s in aio["sources"][:3]:
            print(f"        src: {(s['title'] or '')[:55]!r}")
    for a in it["paidResults"][:3]:
        print(f"    [ad {a['adPosition']}] {a['title'][:44]!r} {(a['url'] or '')[:45]}")
    for p in it["paidProducts"][:3]:
        print(f"    [pla] {p['title'][:40]!r} {p['prices']} {p['displayedUrl']}")
    for p in paa_answered[:3]:
        print(f"    [paa] {p['question'][:48]!r}")
        print(f"          A: {p['answer'][:90]!r}")
        print(f"          src: {p['url'] or '-'} | {(p['title'] or '-')[:45]}")
    assert it["organicResults"], f"{label}: no organic results"
    if expect_ads:
        assert it["paidResults"], f"{label}: expected text ads, got none"
    if expect_products:
        assert it["paidProducts"], f"{label}: expected product ads, got none"
    if expect_paa_answers:
        assert paa_answered, f"{label}: expected PAA answers, got none"
    if expect_sitelinks:
        assert sitelinked, f"{label}: expected sitelinks, got none"
        assert it["suggestedResults"], f"{label}: expected suggestedResults"
    if expect_aio:
        assert aio and aio["content"], f"{label}: expected an AI Overview"
        assert aio["sources"], f"{label}: expected AI Overview sources"
    if expect_device:
        assert it["searchQuery"]["device"] == expect_device, (
            f"{label}: device={it['searchQuery']['device']}"
        )
    if expect_icons:
        iconed = [
            o
            for o in it["organicResults"]
            if (o["icon"] or "").startswith("data:image")
        ]
        print(
            f"    [icons] {len(iconed)}/{len(it['organicResults'])} organic "
            f"carry a base64 favicon"
        )
        assert iconed, f"{label}: expected base64 icons on organic results"


_CASES = {
    "plain": lambda: run("plain query", queries="python asyncio tutorial"),
    "site": lambda: run("site: filter", queries="machine learning", site="arxiv.org"),
    "ads": lambda: run("text ads", queries="car insurance quotes", expect_ads=True),
    "products": lambda: run(
        "product ads", queries="buy running shoes", expect_products=True
    ),
    "focus": lambda: run(
        "focusOnPaidAds (commercial)",
        queries="car insurance quotes",
        focusOnPaidAds=True,
        expect_ads=True,
    ),
    "focus-neg": lambda: run(
        "focusOnPaidAds (non-commercial, retries capped)",
        queries="python asyncio tutorial",
        focusOnPaidAds=True,
    ),
    "paa": lambda: run(
        "people also ask", queries="what is seo", expect_paa_answers=True
    ),
    "sitelinks": lambda: run(
        "sitelinks + suggested (brand query)", queries="amazon", expect_sitelinks=True
    ),
    "aio": lambda: run("AI Overview", queries="benefits of green tea", expect_aio=True),
    "mobile": lambda: run(
        "mobile layout (mobileResults)",
        queries="best seo tools",
        mobileResults=True,
        expect_device="MOBILE",
    ),
    "unfiltered": lambda: run(
        "includeUnfilteredResults (filter=0)",
        queries="python asyncio tutorial",
        includeUnfilteredResults=True,
    ),
    "icons": lambda: run(
        "includeIcons (base64 favicons)",
        queries="github",
        includeIcons=True,
        expect_icons=True,
    ),
    "aimode": lambda: run_ai_mode(
        "Google AI Mode (udm=50)", queries="what is quantum computing"
    ),
}


async def main() -> None:
    names = sys.argv[1:] or list(_CASES)
    try:
        for name in names:
            await _CASES[name]()
    finally:
        await close_sessions()
    print("\nALL E2E OK")


if __name__ == "__main__":
    asyncio.run(main())