werkrevisor/tests/fakes.py

import copy
import io
import os
import tempfile
import tomllib

from werkrevisor.config import PullRequestRef, parse_config
from werkrevisor.review import ModelRun

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
HEAD = "a" * 40
BASE = "b" * 40
BOT = "werkrevisor-bot"


def config(**overrides):
    """The werkprobe example config with per-section overrides, e.g. gate={"yellow_gate": ...}."""
    with open(os.path.join(ROOT, "config", "werkprobe.toml"), "rb") as f:
        data = tomllib.load(f)
    for section, values in overrides.items():
        data.setdefault(section, {}).update(values)
    return parse_config(data)


def pr_ref(head=HEAD):
    return PullRequestRef(base_url="https://git.example.org", owner="team", repo="app", number=7,
                          head_sha=head, base_sha=BASE, base_ref="main", head_ref="feature")


def gitea_file(path, status="modified", additions=1, deletions=0):
    return {"filename": path, "status": status, "additions": additions, "deletions": deletions}


def diff_for(path, added=(), removed=(), new=False, deleted=False):
    lines = [f"diff --git a/{path} b/{path}"]
    if new:
        lines.append("new file mode 100644")
    if deleted:
        lines.append("deleted file mode 100644")
    lines += [f"--- {'/dev/null' if new else 'a/' + path}", f"+++ {'/dev/null' if deleted else 'b/' + path}"]
    lines.append(f"@@ -1,{len(removed)} +1,{len(added)} @@")
    lines += ["-" + r for r in removed] + ["+" + a for a in added]
    return "\n".join(lines) + "\n"


class FakeGitea:
    def __init__(self, files=(), diff="", head=HEAD, state="open", comments=(), reviews=(), author="dev"):
        self.files = list(files)
        self.diff = diff
        self.heads = [head]
        self.state = state
        self.comments = [copy.deepcopy(c) for c in comments]
        self.reviews = list(reviews)
        self.requested = []
        self.statuses = []
        self.created_comments = []
        self.writes = []
        self.author = author
        self.created_reviews = []
        self.fail_review_request = False
        self.raw_files = {}
        self.raw_requests = []

    def get_pr(self, number):
        head = self.heads.pop(0) if len(self.heads) > 1 else self.heads[0]
        return {"number": number, "state": self.state, "title": "Title", "body": "Body",
                "head": {"sha": head}, "base": {"sha": BASE}, "user": {"login": self.author},
                "requested_reviewers": [{"login": login} for login in self.requested]}

    def current_user(self):
        return {"login": BOT}

    def list_pr_files(self, number):
        return self.files

    def get_pr_diff(self, number):
        return self.diff

    def get_raw_file(self, path, ref):
        self.raw_requests.append((path, ref))
        return self.raw_files.get(path, b"line 1\nline 2\n")

    def list_issue_comments(self, number):
        return self.comments

    def create_issue_comment(self, number, body):
        self.writes.append("comment")
        self.created_comments.append(body)
        url = f"https://git.example.org/team/app/pulls/7#issuecomment-{len(self.comments) + 1}"
        self.comments.append({"user": {"login": BOT}, "body": body, "html_url": url})
        return {"html_url": url}

    def set_status(self, sha, state, context, description, target_url=""):
        self.writes.append("status")
        self.statuses.append((sha, state, context, description))

    def list_reviews(self, number):
        return self.reviews

    def request_reviewers(self, number, logins):
        if self.fail_review_request:
            raise OSError("HTTP 422")
        self.writes.append("request")
        self.requested.extend(logins)

    def create_review(self, number, event, body, commit_id):
        if self.author == BOT:
            raise OSError("HTTP 422: approve your own pull is not allowed")
        self.writes.append("review")
        self.created_reviews.append((event, commit_id, body))
        self.reviews.append({"state": event, "commit_id": commit_id, "user": {"login": BOT}, "body": body})


class FakeKey:
    def __init__(self, remaining=10.0, fail=False):
        self.remaining = remaining
        self.fail = fail
        self.calls = 0

    def fetch(self):
        self.calls += 1
        if self.fail:
            raise OSError("network down")
        return {"limit_remaining": self.remaining, "usage": 1.0}


def review_data(recommendation="safe_to_merge", risk="low", issues=(), security="No"):
    return {"review": {
        "merge_recommendation": recommendation,
        "risk_level": risk,
        "security_concerns": security,
        "key_issues_to_review": [
            {"relevant_file": f, "issue_header": "Possible Bug", "issue_content": "Breaks X.",
             "start_line": 3, "end_line": 5} for f in issues],
    }}


class FakeModel:
    def __init__(self, data=None, prompt_files=None, head=HEAD, details=None, error=None, **run_kwargs):
        self.data = data if data is not None else review_data()
        self.prompt_files = prompt_files
        self.head = head
        self.details = details or {"num_ai_calls": 1, "prompt_tokens": 1000, "completion_tokens": 300,
                                   "total_cost_usd": 0.001, "cost_status": "complete"}
        self.error = error
        self.run_kwargs = run_kwargs
        self.calls = []
        self.sections = []

    def __call__(self, config, pr, instructions, full_files_section=""):
        self.calls.append(instructions)
        self.sections.append(full_files_section)
        if self.error:
            raise self.error
        return ModelRun(data=self.data, details=self.details, prompt_files=set(self.prompt_files or ()),
                        reviewed_head_sha=self.head, **self.run_kwargs)


def deps(gitea, model=None, key=None, jev=None):
    from werkrevisor.cli import Deps
    usage_dir = tempfile.mkdtemp(prefix="werkrevisor-test-")
    return Deps(gitea=gitea, key_client=key or FakeKey(), model_review=model or FakeModel(),
                is_reviewable=lambda path: not path.endswith(".png"), out=io.StringIO(),
                usage_path=os.path.join(usage_dir, "usage.jsonl"), jev=jev)


def jev_answer(question, problem, confidence=0.9):
    """A Jev answer whose problem value is `problem`, in the shape of the given question's type."""
    if question["type"] == "noul":
        return {"type": "noul", "noul": problem}
    if question["type"] == "score":
        top = len(question["criteria"]) - 1
        return {"type": "score", "score": problem * top, "confidence": confidence,
                "probabilities": {str(i): 1.0 / (top + 1) for i in range(top + 1)}}
    raise ValueError("choice answers need explicit probabilities")


class FakeJev:
    """problems maps a question name, or (scenario id or section title, question name), to a problem value."""

    def __init__(self, problems=None, confidence=0.9, error=None, model="jev-test"):
        self.problems = problems or {}
        self.confidence = confidence
        self.error = error
        self.model = model
        self.requests = []

    def ask(self, state, questions):
        from werkrevisor.jev import JevAnswer
        self.requests.append((state, questions))
        if self.error:
            raise self.error
        unit = (state.get("scenario") or {}).get("id") or (state.get("section") or {}).get("title")
        answers = {}
        for name, question in questions.items():
            problem = self.problems.get((unit, name), self.problems.get(name, 0.0))
            answers[name] = jev_answer(question, problem, self.confidence)
        return JevAnswer(model=self.model, answers=answers, input_tokens=1000, output_tokens=10, attempts=1)