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):
        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"
    keyfile = tmp / "key"
    raw = secrets.token_urlsafe(32)
    keyfile.write_text(raw + "\n")
    keyfile.chmod(0o600)
    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()}"
    proc = subprocess.Popen(
        [
            str(BIN),
            "serve",
            "-listen",
            listen,
            "-keys",
            str(keys),
            "-test-allow-prefix",
            "127.0.0.0/8",
        ],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    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
        md = env["data"]["markdown"]
        assert "hello" in md.lower() or "contract" in md.lower(), md
        assert env["data"]["metadata"]["title"]

        # 403 guard
        req = urllib.request.Request(
            f"http://{listen}/v2/scrape",
            data=json.dumps({"url": "http://10.0.0.5/"}).encode(),
            headers={
                "Authorization": f"Bearer {raw}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        try:
            urllib.request.urlopen(req, timeout=10)
            print("expected 403", file=sys.stderr)
            return 1
        except urllib.error.HTTPError as e:
            if e.code != 403:
                print("want 403 got", e.code, file=sys.stderr)
                return 1

        # Real SDK if installed.
        try:
            from firecrawl import Firecrawl  # type: ignore
        except ImportError:
            print("firecrawl-py not installed; JSON contract passed")
            return 0

        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
        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())