werkrevisor/tests/test_pr_lens_script.py

"""Runs werkrevisor/pr_lens.mjs with stub modules of PR Lens against a local stand-in of OpenRouter."""
import http.server
import json
import os
import shutil
import subprocess
import tempfile
import threading
import time
import unittest

from werkrevisor import prlens

NODE = shutil.which("node") or (prlens.NODE if os.path.exists(prlens.NODE) else None)

STUBS = {
    "@coldtea/pr-lens-cli/dist/prompt.js":
        "export const SYSTEM_PROMPT = 's';\nexport const buildExtractionPrompt = () => 'u';\n",
    "@coldtea/pr-lens-cli/dist/skill-content.generated.js": "export const GRAPH_DOCUMENT_JSON_SCHEMA = {};\n",
    "@coldtea/pr-lens-cli/dist/version.js": "export const CLI_VERSION = '0'; export const GENERATOR_NAME = 'g';\n",
    # Two calls, as with the repair round of PR Lens; the first answer is unusable.
    "@coldtea/pr-lens-cli/dist/extract.js": """
export const extractGraph = async (options, complete) => {
  await complete({});
  const text = await complete({});
  if (process.env.STUB_FAIL) throw Object.assign(new Error("invalid"), { code: "MODEL_OUTPUT_INVALID" });
  return { document: { title: text, nodes: [], edges: [], flows: [] }, attempts: 2 };
};
""",
    "@coldtea/pr-lens-cli/dist/providers/index.js": """
export const resolveProvider = (options) => ({ ...options });
export const completeJson = async (provider, call) => {
  const response = await fetch(`${provider.baseUrl}/chat/completions`,
    { method: "POST", headers: { "content-type": "application/json" },
      body: JSON.stringify({ model: provider.model }) });
  return (await response.json()).choices[0].message.content;
};
""",
    "@coldtea/pr-lens-renderer/dist/index.js": "export const renderAll = () => ({ assets: [] });\n",
}


class OpenRouter(http.server.BaseHTTPRequestHandler):
    requests = []
    costs = [0.01, 0.02]
    delay = 0

    def do_POST(self):
        time.sleep(OpenRouter.delay)
        body = json.loads(self.rfile.read(int(self.headers["content-length"])))
        OpenRouter.requests.append(body)
        usage = {"prompt_tokens": 100, "completion_tokens": 10}
        cost = OpenRouter.costs[(len(OpenRouter.requests) - 1) % len(OpenRouter.costs)]
        if cost is not None:
            usage["cost"] = cost
        answer = json.dumps({"choices": [{"message": {"content": "{}"}, "finish_reason": "stop"}], "usage": usage})
        self.send_response(200)
        self.send_header("content-type", "application/json")
        self.end_headers()
        self.wfile.write(answer.encode())

    def log_message(self, *args):
        pass


@unittest.skipUnless(NODE, "Node.js is not installed")
class PrLensScriptTest(unittest.TestCase):
    def setUp(self):
        self.modules = tempfile.mkdtemp()
        self.addCleanup(shutil.rmtree, self.modules)
        for path, text in STUBS.items():
            os.makedirs(os.path.dirname(os.path.join(self.modules, path)), exist_ok=True)
            with open(os.path.join(self.modules, path), "w") as f:
                f.write(text)
        OpenRouter.requests, OpenRouter.costs, OpenRouter.delay = [], [0.01, 0.02], 0
        self.server = http.server.HTTPServer(("127.0.0.1", 0), OpenRouter)
        threading.Thread(target=self.server.serve_forever, daemon=True).start()
        self.addCleanup(self.server.server_close)
        self.addCleanup(self.server.shutdown)

    def request(self):
        return {"patch": "diff", "maxDiffBytes": 1000, "files": [], "additions": 0, "deletions": 0,
                "model": "a/b", "baseUrl": f"http://127.0.0.1:{self.server.server_port}", "apiKeyEnv": "KEY",
                "repo": {"owner": "o", "name": "r"}, "host": "h", "number": 1, "url": "u",
                "baseSha": "b", "headSha": "h", "maxOutputTokens": 10}

    def run_script(self, **env):
        completed = subprocess.run([NODE, prlens.SCRIPT, self.modules], input=json.dumps(self.request()),
                                   capture_output=True, text=True, timeout=60,
                                   env={"PATH": os.environ.get("PATH", ""), "KEY": "k", **env})
        lines = [line for line in completed.stderr.splitlines() if line.startswith(prlens.USAGE_LINE)]
        return completed, {"usage": json.loads(lines[-1][len(prlens.USAGE_LINE):])}

    def test_the_script_sums_the_usage_and_asks_openrouter_for_the_cost(self):
        """Scenario#140.01"""
        completed, result = self.run_script()

        self.assertEqual(completed.returncode, 0, completed.stderr)
        usage = result["usage"]
        self.assertEqual((usage["calls"], usage["input_tokens"], usage["output_tokens"]), (2, 200, 20))
        self.assertAlmostEqual(usage["cost_usd"], 0.03)
        self.assertEqual([r.get("usage") for r in OpenRouter.requests], [{"include": True}] * 2)

    def test_a_failed_drawing_still_names_its_usage(self):
        """Scenario#140.02"""
        completed, result = self.run_script(STUB_FAIL="1")

        self.assertEqual(completed.returncode, 1)
        self.assertIn("pr-lens: MODEL_OUTPUT_INVALID: invalid", completed.stderr)
        self.assertEqual(result["usage"]["calls"], 2)
        self.assertAlmostEqual(result["usage"]["cost_usd"], 0.03)

    def test_a_call_without_an_answer_counts_with_an_unknown_cost(self):
        self.server.shutdown()
        self.server.server_close()

        completed, result = self.run_script()

        self.assertEqual(completed.returncode, 1)
        self.assertEqual((result["usage"]["calls"], result["usage"]["cost_usd"]), (1, None))

    def test_a_call_in_flight_counts_with_an_unknown_cost(self):
        OpenRouter.delay = 5
        request = self.request()

        node = subprocess.Popen([NODE, prlens.SCRIPT, self.modules], stdin=subprocess.PIPE, stderr=subprocess.PIPE,
                                stdout=subprocess.DEVNULL, text=True,
                                env={"PATH": os.environ.get("PATH", ""), "KEY": "k"})
        node.stdin.write(json.dumps(request))
        node.stdin.close()
        first = node.stderr.readline()
        node.kill()
        node.wait()
        node.stderr.close()

        self.assertEqual(json.loads(first[len(prlens.USAGE_LINE):])["calls"], 1)
        self.assertIsNone(json.loads(first[len(prlens.USAGE_LINE):])["cost_usd"])

    def test_an_answer_without_a_cost_makes_the_cost_unknown(self):
        OpenRouter.costs = [0.01, None]

        _, result = self.run_script()

        self.assertEqual((result["usage"]["calls"], result["usage"]["cost_usd"]), (2, None))


if __name__ == "__main__":
    unittest.main()