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)