werkrevisor/tests/fakes.py

import copy
import io
import os
import tempfile
import tomllib

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
HEAD = "a" * 40
# The head of the default branch; a PR's base may lie behind it or on another branch.
MAIN = "d" * 40
BASE = "b" * 40
BOT = "werkrevisor-bot"
TEST_CONFIG = os.path.join(ROOT, "tests", "fixtures", "test-config.toml")


def config(**overrides):
    """The test base config with per-section overrides, e.g. gate={"yellow_gate": ...}."""
    with open(TEST_CONFIG, "rb") as f:
        data = tomllib.load(f)
    # Most tests check the LLM path; classifier tests choose their engine mode explicitly.
    data["engine"]["mode"] = "llm"
    # Most tests check one review; the readiness and automerge tests switch both on explicitly.
    data["gate"]["require_wip_transition"] = False
    data["gate"]["automerge"] = False
    for section, values in overrides.items():
        data.setdefault(section, {}).update(values)
    return parse_config(data)


# The repository part of these host configs is the .werkrevisor.toml of each repository.
REPOSITORY_CONFIGS = {"werkrevisor.toml": os.path.join(ROOT, ".werkrevisor.toml"),
                      "werkator.toml": os.path.join(ROOT, "tests", "fixtures", "werkator.werkrevisor.toml"),
                      "werkcrappit.toml": os.path.join(ROOT, "tests", "fixtures", "werkcrappit.werkrevisor.toml"),
                      "werkdock.toml": os.path.join(ROOT, "tests", "fixtures", "werkdock.werkrevisor.toml")}


def example_config(name):
    """Some files in config/ hold only the host part; the repository part comes from REPOSITORY_CONFIGS."""
    with open(os.path.join(ROOT, "config", name), "rb") as f:
        host = tomllib.load(f)
    repo = None
    if name in REPOSITORY_CONFIGS:
        with open(REPOSITORY_CONFIGS[name], "rb") as f:
            repo = tomllib.load(f)
    return build_config(host, repo)


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.withdrawn = []
        self.statuses = []
        self.created_comments = []
        self.writes = []
        self.author = author
        self.created_reviews = []
        self.fail_review_request = False
        self.raw_files = {}
        self.edited_comments = []
        self.review_comments = []
        self.line_comments = {}
        self.refuse_line_comments = False
        # A finished PR: Gitea recorded the change from its WIP title.
        self.timeline = [{"type": "change_title", "old_title": "WIP: Title", "new_title": "Title"}]
        self.scheduled_merges = []
        self.cancelled_merges = []
        self.branch = {"protected": True, "required_approvals": 1, "enable_status_check": True,
                       "status_check_contexts": ["werkator", "ai-review/gate"], "commit": {"id": MAIN}}
        self.schedule_merge_status = None
        # Statuses of the next refusals, before schedule_merge_status applies.
        self.schedule_merge_refusals = []
        self.mergeable = True
        # Content per (path, ref); None answers 404, e.g. for a file missing at the base commit.
        self.raw_files_at = {}
        self.merge_base = BASE
        self.base_sha = BASE
        self.raw_requests = []
        # None leaves the repository out of the base, as an incomplete answer of Gitea would.
        self.base_repo = {"private": False}

    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": self.base_sha, **({"repo": self.base_repo} if self.base_repo is not None else {})}, "merge_base": self.merge_base, "user": {"login": self.author}, "mergeable": self.mergeable,
                "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, max_bytes=None):
        self.raw_requests.append((path, ref))
        if (path, ref) in self.raw_files_at:
            content = self.raw_files_at[(path, ref)]
            if content is None:
                from werkrevisor.gitea import GiteaError
                raise GiteaError(f"GET raw/{path} failed with HTTP 404", 404)
            return content[:max_bytes]
        return self.raw_files.get(path, b"line 1\nline 2\n")[:max_bytes]

    def list_timeline(self, number):
        return self.timeline

    def get_repo(self):
        return {"default_merge_style": "squash", "default_delete_branch_after_merge": True, "default_branch": "main"}

    def get_branch(self, name):
        return self.branch

    def schedule_merge(self, number, style, delete_branch):
        if self.schedule_merge_refusals:
            from werkrevisor.gitea import GiteaError
            status = self.schedule_merge_refusals.pop(0)
            raise GiteaError(f"POST merge failed with HTTP {status}", status)
        if self.schedule_merge_status:
            from werkrevisor.gitea import GiteaError
            raise GiteaError(f"POST merge failed with HTTP {self.schedule_merge_status}", self.schedule_merge_status)
        if self.scheduled_merges:
            from werkrevisor.gitea import GiteaError
            raise GiteaError("POST merge failed with HTTP 409", 409)
        self.writes.append("schedule_merge")
        self.scheduled_merges.append((number, style, delete_branch))

    def cancel_scheduled_merge(self, number):
        self.writes.append("cancel_merge")
        self.cancelled_merges.append(number)
        if not self.scheduled_merges:
            from werkrevisor.gitea import GiteaError
            raise GiteaError("DELETE merge failed with HTTP 404", 404)
        self.scheduled_merges.clear()

    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)
        comment_id = max((c.get("id") or 0 for c in self.comments), default=0) + 1
        url = f"https://git.example.org/team/app/pulls/7#issuecomment-{comment_id}"
        self.comments.append({"id": comment_id, "user": {"login": BOT}, "body": body, "html_url": url})
        return {"id": comment_id, "html_url": url}

    def delete_issue_comment(self, comment_id):
        self.writes.append("delete")
        self.comments = [c for c in self.comments if c.get("id") != comment_id]

    def get_issue_comment(self, comment_id):
        comment = copy.deepcopy(next(c for c in self.comments if c.get("id") == comment_id))
        comment.pop("assets", None)
        return comment

    def list_comment_assets(self, comment_id):
        return copy.deepcopy(next(c for c in self.comments if c.get("id") == comment_id).get("assets", []))

    def create_comment_asset(self, comment_id, name, content, content_type):
        self.writes.append("asset")
        comment = next(c for c in self.comments if c.get("id") == comment_id)
        asset = {"name": name, "size": len(content), "content_type": content_type,
                 "browser_download_url": f"https://git.example.org/attachments/{comment_id}-{name}"}
        comment.setdefault("assets", []).append(asset)
        return asset

    def edit_issue_comment(self, comment_id, body):
        self.writes.append("edit")
        self.edited_comments.append((comment_id, body))
        # Gitea edits issue comments and line comments under the same path.
        for c in self.comments + [c for cs in self.line_comments.values() for c in cs]:
            if c.get("id") == comment_id:
                c["body"] = body

    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 withdraw_review_request(self, number, logins):
        if self.fail_review_request:
            raise OSError("HTTP 422")
        self.writes.append("withdraw")
        self.withdrawn.extend(logins)
        self.requested = [r for r in self.requested if r not in logins]

    def create_review(self, number, event, body, commit_id, comments=()):
        if self.author == BOT:
            raise OSError("HTTP 422: approve your own pull is not allowed")
        if comments and self.refuse_line_comments:
            from werkrevisor.gitea import GiteaError
            status = 422 if self.refuse_line_comments is True else self.refuse_line_comments
            raise GiteaError(f"POST reviews failed with HTTP {status}", status)
        self.writes.append("review")
        self.created_reviews.append((event, commit_id, body))
        self.review_comments.append(list(comments))
        review_id = len(self.reviews) + 100
        self.reviews.append({"id": review_id, "state": event, "commit_id": commit_id, "user": {"login": BOT},
                             "body": body})
        self.line_comments[review_id] = [
            {"id": review_id * 100 + n, "user": {"login": BOT}, "path": c["path"], "position": c["new_position"],
             "resolver": None,
             "body": c["body"], "html_url": f"https://git.example.org/team/app/pulls/7/files#issuecomment-{review_id}-{n}"}
            for n, c in enumerate(comments)]

    def list_review_comments(self, number, review_id):
        return self.line_comments.get(review_id, [])


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

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


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


class FakeModel:
    """Builds a prompt like PR-Agent and passes it to admit before each call; prompt_words pads the user prompt.

    A list of prompt_words makes one call per entry, like a retry; a stop after earlier calls carries their usage.
    """

    def __init__(self, data=None, prompt_files=None, head=HEAD, details=None, error=None, prompt_words=1000,
                 **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.prompt_words = prompt_words if isinstance(prompt_words, list) else [prompt_words]
        self.run_kwargs = run_kwargs
        self.calls = []
        self.sections = []
        self.prompts = []
        self.excluded = []

    def __call__(self, config, pr, instructions, full_files_section="", admit=None, excluded_paths=()):
        from werkrevisor.budget import NoModelCall
        system = f"You are PR-Reviewer.\n{instructions}"
        for done, words in enumerate(self.prompt_words):
            user = "## File: 'x'\n" + "word " * words + full_files_section
            self.prompts.append((system, user))
            try:
                if admit is not None:
                    admit(system, user)
            except NoModelCall as stop:
                if done:
                    stop.details = {"num_ai_calls": done, "prompt_tokens": 1000 * done, "completion_tokens": 0,
                                    "cost_status": "unavailable"}
                raise
        self.calls.append(instructions)
        self.sections.append(full_files_section)
        self.excluded.append(tuple(excluded_paths))
        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", ".jar")), 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)}}
    # A choice problem names the chosen option; without one, "none" is chosen if offered.
    options = list(question["criteria"])
    choice = problem if isinstance(problem, str) else ("none" if "none" in options else options[0])
    rest = (1.0 - confidence) / max(1, len(options) - 1)
    return {"type": "choice", "choice": choice, "confidence": confidence,
            "probabilities": {o: confidence if o == choice else rest for o in options}}


class FakeJev:
    """problems and confidences map a question name, or (scenario id, section title or change statement, question name),
    to a value."""

    def __init__(self, problems=None, confidence=0.9, error=None, model="jev-test", confidences=None):
        self.problems = problems or {}
        self.confidence = confidence
        self.confidences = confidences or {}
        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
        # A finding's state is plain text; its unit is the finding's location line.
        if isinstance(state, str):
            unit = next((line.split(": ", 1)[1] for line in state.splitlines()
                         if line.startswith(("Location: ", "Changeset id: ", "Build file: ", "Changed file: "))), None)
        else:
            unit = ((state.get("scenario") or {}).get("id") or (state.get("section") or {}).get("title")
                    or (state.get("change") or {}).get("statement") or state.get("changed_file"))
        answers = {}
        for name, question in questions.items():
            problem = self.problems.get((unit, name), self.problems.get(name, 0.0))
            confidence = self.confidences.get((unit, name), self.confidences.get(name, self.confidence))
            answers[name] = jev_answer(question, problem, confidence)
        return JevAnswer(model=self.model, answers=answers, input_tokens=1000, output_tokens=10, attempts=1)