werkrevisor/tests/test_pr_agent_run.py

"""Runs the real PR-Agent 0.47.0 review code with a fake Gitea provider and a fake model, fully offline."""

import unittest
from types import SimpleNamespace

from tests.fakes import HEAD, config, pr_ref

try:
    import pr_agent.git_providers as providers
    from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
    from pr_agent.algo.run_details import record_ai_call
    from pr_agent.algo.types import EDIT_TYPE, FilePatchInfo
    from pr_agent.git_providers.git_provider import GitProvider, IncrementalPR
    from pr_agent.git_providers.gitea_provider import GiteaProvider, _GiteaCommitAdapter
    HAVE_PR_AGENT = True
except ImportError:
    HAVE_PR_AGENT = False

ANSWER = """```yaml
review:
  estimated_effort_to_review_[1-5]: |
    1
  risk_level: |
    low
  merge_recommendation: |
    changes_required
  relevant_tests: |
    No
  key_issues_to_review:
    - relevant_file: |
        src/a.py
      issue_header: |
        Possible Bug
      issue_content: |
        Division by zero when n is 0.
      start_line: 2
      end_line: 2
  security_concerns: |
    No
```"""

PATCH = "@@ -1,1 +1,2 @@\n def f(n):\n+    return 1 / n\n"


def _patch_info(name, patch):
    return FilePatchInfo(base_file="def f(n):\n", head_file="def f(n):\n    return 1 / n\n", patch=patch,
                         filename=name, edit_type=EDIT_TYPE.MODIFIED, num_plus_lines=1, num_minus_lines=0)


def fake_provider_class(files):
    class FakeGiteaProvider(GiteaProvider):
        def __init__(self, url=None):
            GitProvider.__init__(self)
            self.pr_url = url
            self.base_url = "https://git.example.org"
            self.owner, self.repo, self.pr_number = "team", "app", 7
            self.pr = SimpleNamespace(title="Add f", body="Adds f.", labels=[], html_url=url,
                                      head=SimpleNamespace(sha=HEAD, ref="feature"),
                                      base=SimpleNamespace(sha="b" * 40, ref="main"))
            self.sha = HEAD
            self.git_files = [{"filename": f.filename, "status": "modified"} for f in files]
            self.diff_files = list(files)
            self.incremental = IncrementalPR(False)
            self.last_commit = self.last_commit_id = _GiteaCommitAdapter({"sha": HEAD})
            self.enabled_pr, self.enabled_issue = True, False
            self.max_comment_chars = 65000

        def get_languages(self):
            return {"Python": 100}

        def get_commit_messages(self):
            return "Add f"

        def get_diff_files(self):
            return self.diff_files

        def publish_comment(self, *args, **kwargs):
            raise AssertionError("publish_output is false, nothing may be published")

    return FakeGiteaProvider


def fake_handler_class(answer, prompts):
    class FakeHandler(BaseAiHandler):
        def __init__(self):
            pass

        @property
        def deployment_id(self):
            return None

        async def chat_completion(self, model, system, user, temperature=0.2, img_path=None):
            prompts.append((model, system, user))
            record_ai_call({"prompt_tokens": 1234, "completion_tokens": 210, "total_tokens": 1444}, model=model)
            return answer, "stop"

    return FakeHandler


@unittest.skipUnless(HAVE_PR_AGENT, "runs inside the PR-Agent image")
class PrAgentRunTest(unittest.TestCase):
    def setUp(self):
        from werkrevisor.redact import Redactor
        from werkrevisor.review import configure_logging
        configure_logging(Redactor(), "CRITICAL")

    def run_review(self, files, answer=ANSWER, section="", **config_overrides):
        from werkrevisor.review import run_model_review
        prompts = []
        original = providers._GIT_PROVIDERS["gitea"]
        providers._GIT_PROVIDERS["gitea"] = fake_provider_class(files)
        try:
            run = run_model_review(config(**config_overrides), pr_ref(), "t" * 20, "k" * 20, "Check X.", section,
                                   ai_handler=fake_handler_class(answer, prompts))
        finally:
            providers._GIT_PROVIDERS["gitea"] = original
        return run, prompts

    def test_review_result_usage_and_prompt(self):
        run, prompts = self.run_review([_patch_info("src/a.py", PATCH)])
        self.assertEqual(len(prompts), 1)
        model, system, user = prompts[0]
        self.assertEqual(model, config().model.id)
        self.assertIn("Check X.", system)
        self.assertIn("merge_recommendation", system)
        self.assertIn("## File: 'src/a.py'", user)
        self.assertEqual(run.prompt_files, {"src/a.py"})
        self.assertFalse(run.pruned)
        self.assertEqual(run.remaining_files, [])
        self.assertEqual(run.reviewed_head_sha, HEAD)
        self.assertEqual(run.data["review"]["merge_recommendation"].strip(), "changes_required")
        self.assertEqual(run.details["num_ai_calls"], 1)
        self.assertEqual((run.details["prompt_tokens"], run.details["completion_tokens"]), (1234, 210))
        self.assertEqual(run.details["cost_status"], "unavailable")

    def test_full_files_follow_the_diff_in_the_user_prompt(self):
        from werkrevisor import fullfiles
        section = fullfiles.render_section([fullfiles.render_file("src/a.py", "def f(n):\n    return 1 / n\n")], HEAD)
        run, prompts = self.run_review([_patch_info("src/a.py", PATCH)], section=section)
        _, system, user = prompts[0]
        self.assertNotIn("Full content of the changed files", system)
        diff_at = user.index("## File: 'src/a.py'")
        files_at = user.index("## Full content of the changed files")
        self.assertLess(diff_at, files_at)
        self.assertLess(files_at, user.index(fullfiles.SECTION_END))
        self.assertEqual(run.prompt_files, {"src/a.py"})

    def test_pruning_is_detected(self):
        big = "@@ -1,1 +1,4001 @@\n def f(n):\n" + "".join(f"+    x{i} = {i} * n\n" for i in range(4000))
        files = [_patch_info("src/a.py", PATCH), _patch_info("src/big.py", big)]
        run, _ = self.run_review(files, model={"max_input_tokens": 3000})
        self.assertTrue(run.pruned)
        self.assertEqual(run.remaining_files, ["src/big.py"])
        self.assertEqual(run.prompt_files, {"src/a.py"})


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