| 1 | """GitHubClient accounting and budgeter pacing — fake transports only."""
|
| 2 |
|
| 3 | from __future__ import annotations
|
| 4 |
|
| 5 | import json
|
| 6 | import sys
|
| 7 | import unittest
|
| 8 | from pathlib import Path
|
| 9 |
|
| 10 | sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts"))
|
| 11 |
|
| 12 | from oakbench.platform_api import ( # noqa: E402
|
| 13 | GitHubApiError,
|
| 14 | GitHubClient,
|
| 15 | RateLimitBudgeter,
|
| 16 | )
|
| 17 |
|
| 18 |
|
| 19 | class FakeTransport:
|
| 20 | """Canned-response transport; records every request, never the network."""
|
| 21 |
|
| 22 | def __init__(self, responses=None) -> None:
|
| 23 | self.requests: list[dict] = []
|
| 24 | self.responses = list(responses or [])
|
| 25 |
|
| 26 | def __call__(self, request: dict) -> dict:
|
| 27 | self.requests.append(request)
|
| 28 | if self.responses:
|
| 29 | return self.responses.pop(0)
|
| 30 | return ok({"ok": True})
|
| 31 |
|
| 32 |
|
| 33 | def ok(payload, *, remaining="4999", request_id="REQ-1", status=200) -> dict:
|
| 34 | return {
|
| 35 | "status": status,
|
| 36 | "headers": {
|
| 37 | "X-RateLimit-Remaining": remaining,
|
| 38 | "X-GitHub-Request-Id": request_id,
|
| 39 | },
|
| 40 | "body": json.dumps(payload).encode("utf-8"),
|
| 41 | }
|
| 42 |
|
| 43 |
|
| 44 | class GitHubClientMetricsTest(unittest.TestCase):
|
| 45 | def test_metrics_accumulate_across_calls(self) -> None:
|
| 46 | first = ok({"name": "repo"}, remaining="4999", request_id="A")
|
| 47 | second = ok({"object": {"sha": "abc123"}}, remaining="4998", request_id="B")
|
| 48 | transport = FakeTransport([first, second])
|
| 49 | client = GitHubClient(token="t0ken", transport=transport)
|
| 50 |
|
| 51 | client.get_repo("acme/widgets")
|
| 52 | sha = client.get_main_sha("acme/widgets")
|
| 53 |
|
| 54 | self.assertEqual(sha, "abc123")
|
| 55 | metrics = client.api_metrics
|
| 56 | self.assertEqual(metrics["request_count"], 2)
|
| 57 | self.assertEqual(
|
| 58 | metrics["response_bytes"],
|
| 59 | len(first["body"]) + len(second["body"]),
|
| 60 | )
|
| 61 | self.assertGreater(metrics["request_bytes"], 0)
|
| 62 | self.assertEqual(metrics["rate_limit_remaining"], 4998)
|
| 63 | self.assertEqual(metrics["request_ids"], ["A", "B"])
|
| 64 |
|
| 65 | def test_request_bytes_include_post_body(self) -> None:
|
| 66 | transport = FakeTransport([ok({"number": 7}), ok({"number": 8})])
|
| 67 | client = GitHubClient(token="t", transport=transport)
|
| 68 | client.create_pr("acme/widgets", title="t", head="b", base="main")
|
| 69 | bytes_with_body = client.api_metrics["request_bytes"]
|
| 70 |
|
| 71 | bare = GitHubClient(token="t", transport=FakeTransport([ok({})]))
|
| 72 | bare.get_repo("acme/widgets")
|
| 73 | self.assertGreater(bytes_with_body, 0)
|
| 74 | sent = transport.requests[0]
|
| 75 | self.assertEqual(sent["method"], "POST")
|
| 76 | self.assertIn(b'"head": "b"', sent["body"])
|
| 77 |
|
| 78 | def test_rate_limit_header_is_case_insensitive_and_last_seen_wins(self) -> None:
|
| 79 | responses = [
|
| 80 | {
|
| 81 | "status": 200,
|
| 82 | "headers": {"x-ratelimit-remaining": "10", "x-github-request-id": "low"},
|
| 83 | "body": b"{}",
|
| 84 | },
|
| 85 | {
|
| 86 | "status": 200,
|
| 87 | "headers": {"X-RATELIMIT-REMAINING": "9"},
|
| 88 | "body": b"{}",
|
| 89 | },
|
| 90 | ]
|
| 91 | client = GitHubClient(token="t", transport=FakeTransport(responses))
|
| 92 | client.get_repo("a/b")
|
| 93 | client.get_repo("a/b")
|
| 94 | self.assertEqual(client.api_metrics["rate_limit_remaining"], 9)
|
| 95 | self.assertEqual(client.api_metrics["request_ids"], ["low"])
|
| 96 |
|
| 97 | def test_request_ids_capped_at_50_keeping_most_recent(self) -> None:
|
| 98 | responses = [ok({}, request_id=f"id-{i}") for i in range(55)]
|
| 99 | client = GitHubClient(token="t", transport=FakeTransport(responses))
|
| 100 | for _ in range(55):
|
| 101 | client.get_repo("a/b")
|
| 102 | ids = client.api_metrics["request_ids"]
|
| 103 | self.assertEqual(len(ids), 50)
|
| 104 | self.assertEqual(ids[0], "id-5")
|
| 105 | self.assertEqual(ids[-1], "id-54")
|
| 106 | self.assertEqual(client.api_metrics["request_count"], 55)
|
| 107 |
|
| 108 | def test_error_status_raises_but_still_counts(self) -> None:
|
| 109 | response = ok({"message": "Validation Failed"}, status=422, remaining="4000")
|
| 110 | client = GitHubClient(token="t", transport=FakeTransport([response]))
|
| 111 | with self.assertRaises(GitHubApiError) as ctx:
|
| 112 | client.create_ref("a/b", "refs/heads/x", "deadbeef")
|
| 113 | self.assertEqual(ctx.exception.status, 422)
|
| 114 | self.assertEqual(client.api_metrics["request_count"], 1)
|
| 115 | self.assertEqual(client.api_metrics["rate_limit_remaining"], 4000)
|
| 116 |
|
| 117 | def test_auth_header_present_and_paths_correct(self) -> None:
|
| 118 | transport = FakeTransport(
|
| 119 | [ok({}), ok({}), ok({}), ok({}), {"status": 204, "headers": {}, "body": b""}]
|
| 120 | )
|
| 121 | client = GitHubClient(token="secret", transport=transport)
|
| 122 | client.put_file("a/b", path="d/f.md", message="m", content=b"hi", branch="x")
|
| 123 | client.get_pr("a/b", 12)
|
| 124 | client.merge_pr("a/b", 12)
|
| 125 | client.create_ref("a/b", "topic", "sha1")
|
| 126 | client.delete_ref("a/b", "refs/heads/topic")
|
| 127 |
|
| 128 | urls = [request["url"] for request in transport.requests]
|
| 129 | self.assertEqual(
|
| 130 | urls,
|
| 131 | [
|
| 132 | "https://api.github.com/repos/a/b/contents/d/f.md",
|
| 133 | "https://api.github.com/repos/a/b/pulls/12",
|
| 134 | "https://api.github.com/repos/a/b/pulls/12/merge",
|
| 135 | "https://api.github.com/repos/a/b/git/refs",
|
| 136 | "https://api.github.com/repos/a/b/git/refs/heads/topic",
|
| 137 | ],
|
| 138 | )
|
| 139 | for request in transport.requests:
|
| 140 | self.assertEqual(request["headers"]["Authorization"], "Bearer secret")
|
| 141 | # bare branch name normalized to a fully qualified ref
|
| 142 | ref_payload = json.loads(transport.requests[3]["body"].decode())
|
| 143 | self.assertEqual(ref_payload["ref"], "refs/heads/topic")
|
| 144 |
|
| 145 |
|
| 146 | class RateLimitBudgeterTest(unittest.TestCase):
|
| 147 | def test_first_call_needs_no_wait(self) -> None:
|
| 148 | clock = FakeBudgetClock(0.0)
|
| 149 | budgeter = RateLimitBudgeter(ops_per_hour=1800, clock=clock)
|
| 150 | self.assertEqual(budgeter.pace(), 0.0)
|
| 151 | self.assertEqual(budgeter.injected_wait_s_total, 0.0)
|
| 152 |
|
| 153 | def test_back_to_back_calls_pay_full_spacing(self) -> None:
|
| 154 | clock = FakeBudgetClock(0.0)
|
| 155 | budgeter = RateLimitBudgeter(ops_per_hour=1800, clock=clock) # 2s spacing
|
| 156 | budgeter.pace()
|
| 157 | self.assertEqual(budgeter.pace(), 2.0)
|
| 158 | self.assertEqual(budgeter.injected_wait_s_total, 2.0)
|
| 159 |
|
| 160 | def test_elapsed_time_reduces_the_injected_wait(self) -> None:
|
| 161 | clock = FakeBudgetClock(0.0)
|
| 162 | budgeter = RateLimitBudgeter(ops_per_hour=1800, clock=clock)
|
| 163 | budgeter.pace()
|
| 164 | clock.now = 1.5
|
| 165 | self.assertEqual(budgeter.pace(), 0.5)
|
| 166 | clock.now = 10.0 # far past the next slot: no wait needed
|
| 167 | self.assertEqual(budgeter.pace(), 0.0)
|
| 168 | self.assertEqual(budgeter.injected_wait_s_total, 0.5)
|
| 169 |
|
| 170 | def test_rejects_nonpositive_budget(self) -> None:
|
| 171 | with self.assertRaises(ValueError):
|
| 172 | RateLimitBudgeter(ops_per_hour=0)
|
| 173 |
|
| 174 |
|
| 175 | class FakeBudgetClock:
|
| 176 | def __init__(self, now: float) -> None:
|
| 177 | self.now = now
|
| 178 |
|
| 179 | def __call__(self) -> float:
|
| 180 | return self.now
|
| 181 |
|
| 182 |
|
| 183 | if __name__ == "__main__":
|
| 184 | unittest.main()
|