devplacepy/tests/e2e/nodes/extension.py

# retoor <retoor@molodetz.nl>

import json
import os
import shutil
import socket
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

import pytest

from devplacepy.database import get_table, init_db, set_consent
from devplacepy.database.nodes import create_device, save_node_settings, suspend_device, unsuspend_device
from devplacepy.services.nodes.broker import JobBroker
from devplacepy.services.nodes.policy import DestinationPolicy
from devplacepy.utils import register_account
from tests.conftest import PROJECT_ROOT

EXTENSION_SOURCE = Path(PROJECT_ROOT) / "nodes" / "extension"
JOB_URL = "https://www.reddit.com/r/python.json"


def free_port() -> int:
    with socket.socket() as probe:
        probe.bind(("127.0.0.1", 0))
        return probe.getsockname()[1]


@pytest.fixture(scope="module")
def donor():
    init_db()
    name = f"extdonor{int(time.time() * 1000)}"
    uid, _, _ = register_account(name, f"{name}@t.dev", "secret123", age_band="adult", accepted_terms=True)
    set_consent("user", uid, "bandwidth_sharing", True)
    save_node_settings(uid, {"enabled": True})
    device, secret = create_device(uid, "Browser under test", "extension")
    return {"uid": uid, "device": device, "secret": secret}


@pytest.fixture(scope="module")
def node_server(donor):
    port = free_port()
    log = open(f"/tmp/node_app_{port}.log", "w")
    proc = subprocess.Popen(
        [sys.executable, "-m", "uvicorn", "--factory", "tests.e2e.nodes.node_app:build", "--host", "127.0.0.1", "--port", str(port), "--log-level", "warning"],
        cwd=str(PROJECT_ROOT),
        env={**os.environ, "PYTHONUNBUFFERED": "1"},
        stdout=log,
        stderr=subprocess.STDOUT,
    )
    base = f"http://127.0.0.1:{port}"
    deadline = time.time() + 30
    while time.time() < deadline:
        try:
            urllib.request.urlopen(f"{base}/stats", timeout=2)
            break
        except Exception:
            if proc.poll() is not None:
                raise RuntimeError(f"node app died, see /tmp/node_app_{port}.log")
            time.sleep(0.3)
    else:
        proc.terminate()
        raise RuntimeError("node app did not start")
    yield base
    proc.terminate()
    proc.wait(timeout=10)
    from devplacepy.database import refresh_snapshot, set_setting

    refresh_snapshot()
    set_setting("nodes_enabled", "0")
    set_setting("nodes_allow_unknown_class", "0")


def stats(base: str) -> dict:
    return json.loads(urllib.request.urlopen(f"{base}/stats", timeout=5).read())


def wait_for(predicate, seconds: float = 20.0, message: str = "condition"):
    deadline = time.time() + seconds
    while time.time() < deadline:
        value = predicate()
        if value:
            return value
        time.sleep(0.2)
    raise AssertionError(f"{message} never became true")


@pytest.fixture
def extension_dir(tmp_path):
    target = tmp_path / "extension"
    shutil.copytree(EXTENSION_SOURCE, target, ignore=shutil.ignore_patterns("test", "node_modules", "*.zip"))
    manifest = json.loads((target / "manifest.json").read_text())
    manifest["host_permissions"].append("http://127.0.0.1/*")
    (target / "manifest.json").write_text(json.dumps(manifest))
    return target


@pytest.fixture
def browser_with_extension(playwright_instance, extension_dir, tmp_path):
    context = playwright_instance.chromium.launch_persistent_context(
        str(tmp_path / "profile"),
        headless=False,
        args=[
            "--headless=new",
            f"--disable-extensions-except={extension_dir}",
            f"--load-extension={extension_dir}",
            "--no-sandbox",
        ],
    )
    worker = context.service_workers[0] if context.service_workers else context.wait_for_event("serviceworker", timeout=20000)
    extension_id = worker.url.split("/")[2]
    yield context, extension_id
    context.close()


def command(page, message: dict) -> dict:
    return page.evaluate("(message) => chrome.runtime.sendMessage(message)", message)


def popup(context, extension_id):
    page = context.new_page()
    page.goto(f"chrome-extension://{extension_id}/ui/popup.html", wait_until="domcontentloaded")
    return page


def test_the_extension_connects_serves_a_job_gets_suspended_and_disconnects(node_server, donor, browser_with_extension):
    context, extension_id = browser_with_extension
    page = popup(context, extension_id)
    assert command(page, {"cmd": "status"})["state"]["state"] == "off"
    page.locator("#disclosure-card").wait_for(state="visible")
    disclosure = context.new_page()
    disclosure.goto(f"chrome-extension://{extension_id}/ui/disclosure.html", wait_until="domcontentloaded")
    disclosure.locator("#accept").click()
    disclosure.locator("#done").filter(has_text="Accepted").wait_for(state="visible")
    reply = command(page, {"cmd": "connect", "serverUrl": node_server, "secret": donor["secret"], "label": "typed locally"})
    assert reply["config"]["hasSecret"] is True and reply["config"]["secretPrefix"].startswith("dpn_")
    assert "dpn_" not in json.dumps(reply)[len('{"config":'):].replace(reply["config"]["secretPrefix"], "")
    assert command(page, {"cmd": "enable", "enabled": True})["config"]["enabled"] is True
    wait_for(lambda: stats(node_server)["connected"] == 1, message="the extension connecting")
    page.reload(wait_until="domcontentloaded")
    page.locator("#state").filter(has_text="Online").wait_for(state="visible", timeout=15000)
    assert page.locator("#label").inner_text() == "Browser under test"
    target = DestinationPolicy().check(JOB_URL, "GET", "fetch")
    uid = JobBroker().enqueue("guest", "browser-e2e", target, "fetch")

    def finished():
        row = get_table("node_requests").find_one(uid=uid)
        return row if row and row["status"] in ("done", "failed") else None

    row = wait_for(finished, seconds=40, message="the job being served")
    assert row["mode"] == "fetch"
    assert row["device_uid"] == donor["device"]["uid"] and row["device_asn_class"] == "unknown"
    assert row["error"] not in ("node_disconnected", "timeout", "send_failed")
    view = command(page, {"cmd": "status"})
    assert view["usage"]["requests"] + view["usage"]["failed"] >= 1
    suspend_device(donor["device"]["uid"], "browser e2e", 30)
    wait_for(lambda: stats(node_server)["connected"] == 0, message="the suspension disconnecting the extension")
    suspended = wait_for(lambda: (state := command(page, {"cmd": "status"})["state"])["state"] == "suspended" and state, message="the suspended state")
    assert suspended["suspendedUntil"].startswith("20")
    page.reload(wait_until="domcontentloaded")
    page.locator("#detail").filter(has_text="Suspended until").wait_for(state="visible", timeout=10000)
    unsuspend_device(donor["device"]["uid"])
    assert command(page, {"cmd": "enable", "enabled": False})["config"]["enabled"] is False
    wait_for(lambda: stats(node_server)["connected"] == 0, message="the switch disconnecting")
    assert command(page, {"cmd": "status"})["state"]["state"] == "off"