werkrevisor/tests/test_config_budget_redact.py

import math
import os
import tempfile
import unittest

from tests.fakes import ROOT, config
from werkrevisor import budget
from werkrevisor.config import ConfigError, load_config, pr_ref_from_env, pr_ref_from_url
from werkrevisor.redact import Redactor
from werkrevisor.verdict import Verdict


class ConfigTest(unittest.TestCase):
    def test_examples_load(self):
        for name, gate in (("werkprobe.toml", "approval"), ("hsadmin-ng.toml", "block")):
            with self.subTest(name=name):
                cfg = load_config(os.path.join(ROOT, "config", name))
                self.assertEqual(cfg.gate.yellow_gate, gate)

    def test_werkprobe_worst_case_fits_the_budget(self):
        cfg = load_config(os.path.join(ROOT, "config", "werkprobe.toml"))
        m = cfg.model
        worst = m.attempts * (m.max_input_tokens * m.input_price_usd_per_mtok
                              + m.max_output_tokens * m.output_price_usd_per_mtok) / 1e6
        self.assertLessEqual(worst, cfg.budget.max_usd_per_review)

    def test_hsadmin_ng_budget(self):
        self.assertEqual(load_config(os.path.join(ROOT, "config", "hsadmin-ng.toml")).budget.max_usd_per_review, 1.0)

    def test_unknown_key_is_rejected(self):
        with self.assertRaisesRegex(ConfigError, "unknown keys: max_finding"):
            config(review={"max_finding": 3})

    def test_invalid_values(self):
        for overrides in ({"gate": {"yellow_gate": "maybe"}},
                          {"gate": {"yellow_gate": "pass_with_review_request", "human_reviewers": []}},
                          {"gate": {"yellow_gate": "approval", "human_reviewers": []}},
                          {"model": {"max_retries": 3}},
                          {"model": {"reasoning_effort": "huge"}},
                          {"categories": {"auto_approvable": ["everything"]}},
                          {"budget": {"max_usd_per_review": 0}}):
            with self.subTest(overrides=overrides):
                with self.assertRaises(ConfigError):
                    config(**overrides)

    def test_missing_file(self):
        with self.assertRaises(ConfigError):
            load_config(os.path.join(tempfile.mkdtemp(), "missing.toml"))

    def test_env(self):
        env = {"WERKATOR_GITEA_BASE_URL": "https://git.example.org/", "WERKATOR_GITEA_OWNER": "o",
               "WERKATOR_GITEA_REPO": "r", "WERKATOR_PR_NUMBER": "12", "WERKATOR_PR_HEAD_SHA": "abc"}
        ref = pr_ref_from_env(env)
        self.assertEqual(ref.url, "https://git.example.org/o/r/pulls/12")
        with self.assertRaisesRegex(ConfigError, "WERKATOR_PR_HEAD_SHA"):
            pr_ref_from_env({k: v for k, v in env.items() if k != "WERKATOR_PR_HEAD_SHA"})
        with self.assertRaisesRegex(ConfigError, "https"):
            pr_ref_from_env(dict(env, WERKATOR_GITEA_BASE_URL="http://git.example.org"))

    def test_pr_url(self):
        ref = pr_ref_from_url("https://git.example.org/gitea/o/r/pulls/5")
        self.assertEqual((ref.base_url, ref.owner, ref.repo, ref.number), ("https://git.example.org/gitea", "o", "r", 5))
        with self.assertRaises(ConfigError):
            pr_ref_from_url("https://git.example.org/o/r/issues/5")


def pinned_config():
    # The example config's prices change with the market; these tests pin their own.
    return config(model={"input_price_usd_per_mtok": 0.30, "output_price_usd_per_mtok": 2.50, "max_output_tokens": 4000},
                     budget={"max_usd_per_review": 0.10})


class BudgetTest(unittest.TestCase):
    def test_estimate_counts_all_attempts(self):
        cfg = pinned_config()
        est = budget.estimate(3500, 0, cfg.model, cfg.budget)
        expected_tokens = math.ceil((3500 * 1.5 + 16000) / 3.5)
        self.assertEqual(est.input_tokens, expected_tokens)
        per_attempt = (expected_tokens * 0.30 + 4000 * 2.50) / 1e6
        self.assertAlmostEqual(est.worst_case_usd, 2 * per_attempt, places=6)

    def test_precheck(self):
        cfg = config()
        self.assertIsNone(budget.precheck(budget.Estimate(1000, 0.01), cfg.model, cfg.budget))
        self.assertEqual(budget.precheck(budget.Estimate(30000, 0.01), cfg.model, cfg.budget).code, "too_large_tokens")
        reason = budget.precheck(budget.Estimate(1000, 0.5), cfg.model, cfg.budget)
        self.assertEqual((reason.code, reason.verdict), ("too_large_cost", Verdict.YELLOW))

    def test_key_check(self):
        b = pinned_config().budget
        self.assertIsNone(budget.key_check({"limit_remaining": 5.0}, b))
        self.assertIsNone(budget.key_check({"limit_remaining": None}, b))
        reason = budget.key_check({"limit_remaining": 0.3}, b)
        self.assertEqual((reason.code, reason.verdict), ("budget", Verdict.ERROR))
        self.assertIsNotNone(budget.key_check({"limit_remaining": 0.08},
                                              config(budget={"min_key_remaining_usd": 0, "max_usd_per_review": 0.10}).budget))

    def test_actual_cost(self):
        m = pinned_config().model
        self.assertEqual(budget.actual_cost_usd({"cost_status": "complete", "total_cost_usd": 0.02}, m),
                         (0.02, "reported"))
        self.assertEqual(budget.actual_cost_usd({"cost_status": "unavailable", "prompt_tokens": 1_000_000,
                                                 "completion_tokens": 0}, m), (0.3, "estimated"))
        self.assertEqual(budget.actual_cost_usd({}, m), (None, "unknown"))


class RedactTest(unittest.TestCase):
    def test_known_secrets_and_patterns(self):
        r = Redactor(["gitea-token-123456"])
        text = r("token gitea-token-123456 key sk-or-v1-abcdef1234567890 Authorization: token xyz123")
        self.assertNotIn("gitea-token-123456", text)
        self.assertNotIn("sk-or-v1-abcdef1234567890", text)
        self.assertNotIn("xyz123", text)

    def test_short_values_are_ignored(self):
        self.assertEqual(Redactor(["", "ab"])("ab cd"), "ab cd")


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