diff --git a/webhook_server/libs/github_api.py b/webhook_server/libs/github_api.py index 0e9b5c5f..a5e53b08 100644 --- a/webhook_server/libs/github_api.py +++ b/webhook_server/libs/github_api.py @@ -7,6 +7,7 @@ import shlex import shutil import tempfile +import threading from typing import Any import github @@ -47,6 +48,31 @@ ) +class CountingRequester: + """ + Wrapper around PyGithub Requester to count API calls for a specific instance. + Intercepts request* methods to increment a counter. + """ + + def __init__(self, requester: Any) -> None: + self._requester = requester + self.count = 0 + self._thread_lock = threading.Lock() + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._requester, name) + if name.startswith("request") and callable(attr): + + def wrapper(*args: Any, **kwargs: Any) -> Any: + # Increment counter with thread safety since PyGithub may run in threads. + with self._thread_lock: + self.count += 1 + return attr(*args, **kwargs) + + return wrapper + return attr + + class GithubWebhook: def __init__(self, hook_data: dict[Any, Any], headers: Headers, logger: logging.Logger) -> None: logger.name = "GithubWebhook" @@ -68,6 +94,8 @@ def __init__(self, hook_data: dict[Any, Any], headers: Headers, logger: logging. self.current_pull_request_supported_retest: list[str] = [] self.github_api: github.Github | None = None self.initial_rate_limit_remaining: int | None = None + self.requester_wrapper: CountingRequester | None = None + self.initial_wrapper_count: int = 0 if not self.config.repository_data: raise RepositoryNotFoundInConfigError(f"Repository {self.repository_name} not found in config file") @@ -80,6 +108,22 @@ def __init__(self, hook_data: dict[Any, Any], headers: Headers, logger: logging. if github_api and self.token: self.github_api = github_api + + # Wrap the requester to count API calls per this webhook instance + # This must be done BEFORE creating self.repository so it shares the wrapped requester + # PyGithub stores the requester in _Github__requester (name mangling) + if hasattr(self.github_api, "_Github__requester"): + requester = self.github_api._Github__requester + if isinstance(requester, CountingRequester): + # Already wrapped (shared Github instance), reuse existing wrapper + self.requester_wrapper = requester + else: + self.requester_wrapper = CountingRequester(requester) + self.github_api._Github__requester = self.requester_wrapper + + # Capture initial count for per-webhook delta calculation + self.initial_wrapper_count = self.requester_wrapper.count + # Track initial rate limit for token spend calculation # Note: log_prefix not set yet, so we can't use it in error messages here try: @@ -143,9 +187,24 @@ async def _get_token_metrics(self) -> str: return "" try: + # Use the wrapper count if available (thread-safe per request) + # We skip checking global rate limit to avoid inflating the API call count with an extra call + if self.requester_wrapper: + token_spend = self.requester_wrapper.count - self.initial_wrapper_count + # If rate limit was already low, or if the delta calculation leads to negative + # values due to race conditions (unlikely with correct locking), clamp to 0 + remaining = max(0, self.initial_rate_limit_remaining - token_spend) + + return ( + f"token {self.token[:8]}... {token_spend} API calls " + f"(initial: {self.initial_rate_limit_remaining}, " + f"remaining: {remaining})" + ) + final_rate_limit = await asyncio.to_thread(self.github_api.get_rate_limit) final_remaining = final_rate_limit.rate.remaining + # Fallback to global rate limit calculation (inaccurate under concurrency) # Calculate token spend (handle case where rate limit reset between checks) # If final > initial, rate limit reset occurred, so we can't calculate accurately if final_remaining > self.initial_rate_limit_remaining: @@ -734,9 +793,10 @@ def __del__(self) -> None: clean up their own directories. Only the base clone directory must be removed here to prevent accumulating stale repositories on disk. """ - if self.clone_repo_dir and os.path.exists(self.clone_repo_dir): + clone_repo_dir = getattr(self, "clone_repo_dir", None) + if clone_repo_dir and os.path.exists(clone_repo_dir): try: - shutil.rmtree(self.clone_repo_dir, ignore_errors=True) - self.logger.debug(f"Cleaned up temp directory (in __del__): {self.clone_repo_dir}") + shutil.rmtree(clone_repo_dir, ignore_errors=True) + self.logger.debug(f"Cleaned up temp directory (in __del__): {clone_repo_dir}") except Exception: pass # Ignore errors during cleanup diff --git a/webhook_server/tests/test_api_call_counting.py b/webhook_server/tests/test_api_call_counting.py new file mode 100644 index 00000000..b725b74b --- /dev/null +++ b/webhook_server/tests/test_api_call_counting.py @@ -0,0 +1,94 @@ +import logging +from unittest.mock import MagicMock, patch + +import pytest +from github import Github +from starlette.datastructures import Headers + +from webhook_server.libs.github_api import CountingRequester, GithubWebhook + + +def test_counting_requester(): + mock_requester = MagicMock() + mock_requester.requestJsonAndCheck.return_value = ("header", "data") + + wrapper = CountingRequester(mock_requester) + + # Initial state + assert wrapper.count == 0 + + # Make a call + wrapper.requestJsonAndCheck("GET", "/foo") + assert wrapper.count == 1 + mock_requester.requestJsonAndCheck.assert_called_once_with("GET", "/foo") + + # Access non-request attribute + _ = wrapper.other_method + # Count should not increase + assert wrapper.count == 1 + + # Make another call + wrapper.requestMultipartAndCheck("POST", "/bar") + assert wrapper.count == 2 + + +@pytest.mark.asyncio +async def test_github_webhook_token_metrics_with_counter(): + # Mock dependencies + mock_logger = MagicMock(spec=logging.Logger) + mock_hook_data = {"repository": {"name": "test-repo", "full_name": "owner/test-repo"}} + mock_headers = Headers({"X-GitHub-Delivery": "123", "X-GitHub-Event": "push"}) + + with ( + patch("webhook_server.libs.github_api.Config") as MockConfig, + patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") as mock_get_api, + patch("webhook_server.libs.github_api.get_github_repo_api"), + patch("webhook_server.libs.github_api.get_repository_github_app_api"), + patch("webhook_server.libs.github_api.prepare_log_prefix"), + # Patch this method to avoid calls to get_apis_and_tokes_from_config which isn't mocked + patch("webhook_server.libs.github_api.GithubWebhook.add_api_users_to_auto_verified_and_merged_users"), + ): + # Setup Config + mock_config = MockConfig.return_value + mock_config.repository_data = {"some": "data"} + + # Setup Github API mock + mock_github = MagicMock(spec=Github) + mock_requester = MagicMock() + # PyGithub uses name mangling for private attributes + mock_github._Github__requester = mock_requester + + # Mock rate limit + mock_rate_limit = MagicMock() + mock_rate_limit.rate.remaining = 5000 + mock_github.get_rate_limit.return_value = mock_rate_limit + + mock_get_api.return_value = (mock_github, "token123", "user") + + # Initialize webhook + webhook = GithubWebhook(mock_hook_data, mock_headers, mock_logger) + + # Verify requester was wrapped + assert isinstance(webhook.requester_wrapper, CountingRequester) + assert webhook.github_api._Github__requester == webhook.requester_wrapper + + # Simulate API calls + webhook.github_api._Github__requester.requestJsonAndCheck("GET", "/foo") + webhook.github_api._Github__requester.requestJsonAndCheck("GET", "/bar") + + assert webhook.requester_wrapper.count == 2 + + # Check metrics output + # Mock final rate limit to be lower (simulating usage by others too) + final_rate_limit = MagicMock() + final_rate_limit.rate.remaining = 4000 # Dropped by 1000 globally + + with patch.object(webhook.github_api, "get_rate_limit", return_value=final_rate_limit): + metrics = await webhook._get_token_metrics() + + # Expect metrics to show 2 calls (our local usage), not 1000 (global usage) + assert "2 API calls" in metrics + assert "initial: 5000" in metrics + # When using wrapper, we don't show "final" anymore, but we show "remaining" + # remaining = initial - spend = 5000 - 2 = 4998 + assert "remaining: 4998" in metrics diff --git a/webhook_server/tests/test_github_api_metrics.py b/webhook_server/tests/test_github_api_metrics.py new file mode 100644 index 00000000..043d7bd1 --- /dev/null +++ b/webhook_server/tests/test_github_api_metrics.py @@ -0,0 +1,287 @@ +import threading +from unittest.mock import Mock, patch + +import pytest +from starlette.datastructures import Headers + +from webhook_server.libs.github_api import CountingRequester, GithubWebhook + + +class TestCountingRequester: + def test_counting_requester_init(self): + requester = Mock() + wrapper = CountingRequester(requester) + assert wrapper._requester == requester + assert wrapper.count == 0 + assert isinstance(wrapper._thread_lock, type(threading.Lock())) + + def test_counting_requester_increments(self): + requester = Mock() + # Mock a request method on the requester + requester.requestJsonAndCheck = Mock(return_value="result") + + wrapper = CountingRequester(requester) + + # Call the method through the wrapper + result = wrapper.requestJsonAndCheck("arg", key="value") + + assert result == "result" + assert wrapper.count == 1 + requester.requestJsonAndCheck.assert_called_once_with("arg", key="value") + + def test_counting_requester_ignores_non_request_methods(self): + requester = Mock() + requester.other_method = Mock(return_value="result") + + wrapper = CountingRequester(requester) + + result = wrapper.other_method() + + assert result == "result" + assert wrapper.count == 0 + + def test_counting_requester_thread_safety(self): + requester = Mock() + requester.requestJson = Mock() + wrapper = CountingRequester(requester) + + def make_requests(): + for _ in range(100): + wrapper.requestJson() + + threads = [] + for _ in range(10): + t = threading.Thread(target=make_requests) + threads.append(t) + t.start() + + for t in threads: + t.join() + + assert wrapper.count == 1000 + + +class TestGithubWebhookMetrics: + @pytest.fixture + def minimal_hook_data(self): + return { + "repository": {"name": "test-repo", "full_name": "org/test-repo"}, + "number": 1, + } + + @pytest.fixture + def minimal_headers(self): + return Headers({"X-GitHub-Event": "pull_request", "X-GitHub-Delivery": "abc"}) + + @pytest.fixture + def logger(self): + return Mock() + + @patch("webhook_server.libs.github_api.Config") + @patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") + @patch("webhook_server.libs.github_api.get_github_repo_api") + @patch("webhook_server.libs.github_api.get_repository_github_app_api") + @patch("webhook_server.utils.helpers.get_repository_color_for_log_prefix") + def test_init_wraps_requester_new( + self, + mock_color, + mock_get_app_api, + mock_get_repo_api, + mock_get_api, + mock_config, + minimal_hook_data, + minimal_headers, + logger, + ): + mock_config.return_value.repository = True + mock_config.return_value.repository_local_data.return_value = {} + + mock_github_api = Mock() + mock_requester = Mock() + mock_github_api._Github__requester = mock_requester + mock_get_api.return_value = (mock_github_api, "token", "apiuser") + + mock_get_repo_api.return_value = Mock() + mock_get_app_api.return_value = Mock() + mock_color.return_value = "test-repo" + + gh = GithubWebhook(minimal_hook_data, minimal_headers, logger) + + # Verify wrapper was created and set + assert isinstance(gh.requester_wrapper, CountingRequester) + assert mock_github_api._Github__requester == gh.requester_wrapper + assert gh.initial_wrapper_count == 0 + + @patch("webhook_server.libs.github_api.Config") + @patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") + @patch("webhook_server.libs.github_api.get_github_repo_api") + @patch("webhook_server.libs.github_api.get_repository_github_app_api") + @patch("webhook_server.utils.helpers.get_repository_color_for_log_prefix") + def test_init_reuses_wrapper( + self, + mock_color, + mock_get_app_api, + mock_get_repo_api, + mock_get_api, + mock_config, + minimal_hook_data, + minimal_headers, + logger, + ): + mock_config.return_value.repository = True + mock_config.return_value.repository_local_data.return_value = {} + + mock_github_api = Mock() + # Already wrapped + existing_wrapper = CountingRequester(Mock()) + existing_wrapper.count = 5 + mock_github_api._Github__requester = existing_wrapper + + mock_get_api.return_value = (mock_github_api, "token", "apiuser") + + mock_get_repo_api.return_value = Mock() + mock_get_app_api.return_value = Mock() + mock_color.return_value = "test-repo" + + gh = GithubWebhook(minimal_hook_data, minimal_headers, logger) + + # Verify wrapper was reused + assert gh.requester_wrapper == existing_wrapper + assert gh.initial_wrapper_count == 5 + + @patch("webhook_server.libs.github_api.Config") + @patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") + @patch("webhook_server.libs.github_api.get_github_repo_api") + @patch("webhook_server.libs.github_api.get_repository_github_app_api") + @patch("webhook_server.utils.helpers.get_repository_color_for_log_prefix") + @pytest.mark.asyncio + async def test_get_token_metrics_with_wrapper( + self, + mock_color, + mock_get_app_api, + mock_get_repo_api, + mock_get_api, + mock_config, + minimal_hook_data, + minimal_headers, + logger, + ): + mock_config.return_value.repository = True + mock_config.return_value.repository_local_data.return_value = {} + + mock_github_api = Mock() + mock_github_api.get_rate_limit.return_value.rate.remaining = 5000 + mock_requester = Mock() + mock_github_api._Github__requester = mock_requester + + mock_get_api.return_value = (mock_github_api, "token", "apiuser") + mock_get_repo_api.return_value = Mock() + mock_get_app_api.return_value = Mock() + + gh = GithubWebhook(minimal_hook_data, minimal_headers, logger) + + # Simulate usage + gh.requester_wrapper.count = 10 + + metrics = await gh._get_token_metrics() + + assert "10 API calls" in metrics + assert "initial: 5000" in metrics + assert "remaining: 4990" in metrics + + @patch("webhook_server.libs.github_api.Config") + @patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") + @patch("webhook_server.libs.github_api.get_github_repo_api") + @patch("webhook_server.libs.github_api.get_repository_github_app_api") + @patch("webhook_server.utils.helpers.get_repository_color_for_log_prefix") + @pytest.mark.asyncio + async def test_get_token_metrics_with_reused_wrapper( + self, + mock_color, + mock_get_app_api, + mock_get_repo_api, + mock_get_api, + mock_config, + minimal_hook_data, + minimal_headers, + logger, + ): + mock_config.return_value.repository = True + mock_config.return_value.repository_local_data.return_value = {} + + mock_github_api = Mock() + mock_github_api.get_rate_limit.return_value.rate.remaining = 4995 + + # Reused wrapper started at 5 + existing_wrapper = CountingRequester(Mock()) + existing_wrapper.count = 5 + mock_github_api._Github__requester = existing_wrapper + + mock_get_api.return_value = (mock_github_api, "token", "apiuser") + mock_get_repo_api.return_value = Mock() + mock_get_app_api.return_value = Mock() + + gh = GithubWebhook(minimal_hook_data, minimal_headers, logger) + + # Initial wrapper count should be 5 + assert gh.initial_wrapper_count == 5 + + # Simulate usage (+3 calls) + gh.requester_wrapper.count = 8 + + metrics = await gh._get_token_metrics() + + # Should show 3 calls (8 - 5) + assert "3 API calls" in metrics + assert "initial: 4995" in metrics + assert "remaining: 4992" in metrics + + @patch("webhook_server.libs.github_api.Config") + @patch("webhook_server.libs.github_api.get_api_with_highest_rate_limit") + @patch("webhook_server.libs.github_api.get_github_repo_api") + @patch("webhook_server.libs.github_api.get_repository_github_app_api") + @patch("webhook_server.utils.helpers.get_repository_color_for_log_prefix") + @pytest.mark.asyncio + async def test_get_token_metrics_fallback_reset( + self, + mock_color, + mock_get_app_api, + mock_get_repo_api, + mock_get_api, + mock_config, + minimal_hook_data, + minimal_headers, + logger, + ): + mock_config.return_value.repository = True + mock_config.return_value.repository_local_data.return_value = {} + + mock_github_api = Mock() + # Initial rate limit + mock_github_api.get_rate_limit.return_value.rate.remaining = 100 + # No wrapper (simulate failure to wrap) + del mock_github_api._Github__requester + + mock_get_api.return_value = (mock_github_api, "token", "apiuser") + mock_get_repo_api.return_value = Mock() + mock_get_app_api.return_value = Mock() + + gh = GithubWebhook(minimal_hook_data, minimal_headers, logger) + + # Manually unset wrapper to test fallback + gh.requester_wrapper = None + + # Mock reset happened (final > initial) + mock_github_api.get_rate_limit.return_value.rate.remaining = 5000 + + metrics = await gh._get_token_metrics() + + assert "rate limit reset occurred" in metrics + assert "initial: 100" in metrics + assert "final: 5000" in metrics + + def test_del_safe_missing_attr(self, minimal_hook_data, minimal_headers, logger): + # Create instance but don't init fully to avoid setting attributes + gh = GithubWebhook.__new__(GithubWebhook) + # Should not raise AttributeError + gh.__del__()