sift/tests/contract/sdk_test.py

#!/usr/bin/env python3
"""SDK contract: firecrawl-py 4.17.0 against a local sifttest sift."""
from __future__ import annotations

import hashlib
import json
import os
import secrets
import socket
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
BIN = ROOT / "bin" / "sift-test"


class Handler(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        if self.path.startswith("/forbid"):
            body = b"nope"
            self.send_response(403)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        if self.path.startswith("/big"):
            body = b"x" * (6 * 1024 * 1024)
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        if self.path.startswith("/json"):
            body = b'{"hello":"contract"}'
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return
        body = b"<!doctype html><html><head><title>Contract</title></head><body><article><p>hello contract</p></article></body></html>"
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)


def free_port() -> int:
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    p = s.getsockname()[1]
    s.close()
    return p


def main() -> int:
    if not BIN.exists():
        print("missing", BIN, file=sys.stderr)
        return 1
    fx_port = free_port()
    httpd = HTTPServer(("127.0.0.1", fx_port), Handler)
    threading.Thread(target=httpd.serve_forever, daemon=True).start()

    tmp = Path(tempfile.mkdtemp(prefix="sift-contract-"))
    keys = tmp / "keys.toml"
    raw = secrets.token_urlsafe(32)
    digest = hashlib.sha256(raw.encode()).hexdigest()
    keys.write_text(
        f'[[caller]]\nname = "contract"\nworkspace = "hermes-fleet"\nkey_sha256 = "{digest}"\n'
    )

    listen = f"127.0.0.1:{free_port()}"
    logf = tmp / "sift.log"
    proc = subprocess.Popen(
        [
            str(BIN),
            "serve",
            "-listen",
            listen,
            "-keys",
            str(keys),
            "-test-allow-prefix",
            "127.0.0.0/8",
        ],
        stdout=logf.open("w"),
        stderr=subprocess.STDOUT,
    )
    try:
        for _ in range(50):
            try:
                urllib.request.urlopen(f"http://{listen}/healthz", timeout=0.2)
                break
            except Exception:
                time.sleep(0.05)
        else:
            print("sift did not start", file=sys.stderr)
            return 1

        page = f"http://127.0.0.1:{fx_port}/"
        req = urllib.request.Request(
            f"http://{listen}/v2/scrape",
            data=json.dumps({"url": page, "formats": ["markdown"]}).encode(),
            headers={
                "Authorization": f"Bearer {raw}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=10) as resp:
            env = json.loads(resp.read())
        assert env.get("success") is True, env

        skip = os.environ.get("SIFT_CONTRACT_SKIP_SDK") == "1"
        try:
            from firecrawl import Firecrawl  # type: ignore
        except ImportError:
            if skip:
                print("SIFT_CONTRACT_SKIP_SDK=1; SDK not installed")
                return 0
            print("firecrawl-py==4.17.0 is required for make contract", file=sys.stderr)
            print("install it, or set SIFT_CONTRACT_SKIP_SDK=1", file=sys.stderr)
            return 1

        app = Firecrawl(api_url=f"http://{listen}", api_key=raw)
        doc = app.scrape(page, formats=["markdown"])
        text = getattr(doc, "markdown", None) or (
            doc.get("markdown") if isinstance(doc, dict) else None
        )
        assert text, doc

        def scrape_err(url: str, key: str | None = None):
            c = Firecrawl(api_url=f"http://{listen}", api_key=key if key is not None else raw)
            try:
                c.scrape(url, formats=["markdown"])
                raise AssertionError(f"expected error for {url}")
            except Exception as e:
                return str(e)

        msg = scrape_err("http://10.0.0.5/")
        assert "blocked:" in msg, msg
        msg = scrape_err(f"http://127.0.0.1:{fx_port}/forbid")
        assert "hint: browser" in msg, msg
        msg = scrape_err(f"http://127.0.0.1:{fx_port}/big")
        assert "413" in msg or "too large" in msg, msg
        msg = scrape_err(page, key="wrong")
        assert "401" in msg or "Unauthorized" in msg or "unauthorized" in msg.lower(), msg

        print("firecrawl-py SDK contract passed")
        return 0
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            proc.kill()
        httpd.shutdown()


if __name__ == "__main__":
    raise SystemExit(main())