Per comments - moving logic to audit from githubclient & moving tests

Author: rreddy2

nodestream_github/audit.py

@@ -6,8 +6,10 @@
 """
 
 from collections.abc import AsyncGenerator
+from datetime import UTC, datetime, timedelta
 from typing import Any
 
+from dateutil.relativedelta import relativedelta
 from nodestream.pipeline import Extractor
 
 from .client import GithubRestApiClient
@@ -17,6 +19,77 @@
 logger = get_plugin_logger(__name__)
 
 
+def generate_date_range(lookback_period: dict[str, int]) -> list[str]:
+    """
+    Generate a list of date strings in YYYY-MM-DD format for
+    the given lookback period.
+    """
+    if not lookback_period:
+        return []
+
+    end_date = datetime.now(tz=UTC).date()
+    start_date = (datetime.now(tz=UTC) - relativedelta(**lookback_period)).date()
+
+    delta_days = (end_date - start_date).days + 1
+    return [
+        (start_date + timedelta(days=i)).strftime("%Y-%m-%d") for i in range(delta_days)
+    ]
+
+
+def build_search_phrase(
+    actions: list[str],
+    actors: list[str],
+    exclude_actors: list[str],
+    target_date: str | None = None,
+) -> str:
+    # adding action-based filtering
+    actions_phrase = ""
+    if actions:
+        actions_phrase = " ".join(f"action:{action}" for action in actions)
+
+    # adding date-based filtering for a specific date
+    date_filter = f"created:{target_date}" if target_date else ""
+
+    # adding actor-based filtering
+    actors_phrase = ""
+    if actors:
+        actors_phrase = " ".join(f"actor:{actor}" for actor in actors)
+
+    # adding exclude_actors based filtering
+    exclude_actors_phrase = ""
+    if exclude_actors:
+        exclude_actors_phrase = " ".join(f"-actor:{actor}" for actor in exclude_actors)
+    return " ".join(
+        section
+        for section in [
+            actions_phrase,
+            date_filter,
+            actors_phrase,
+            exclude_actors_phrase,
+        ]
+        if section
+    ).strip()
+
+
+def validate_lookback_period(lookback_period: dict[str, int]) -> dict[str, int]:
+    """Sanitize the lookback period to only include valid keys."""
+
+    def validate_positive_int(value: int) -> int:
+        converted = int(value)
+        if converted <= 0:
+            negative_value_exception_msg = (
+                f"Lookback period values must be positive: {value}"
+            )
+            raise ValueError(negative_value_exception_msg)
+        return converted
+
+    try:
+        return {k: validate_positive_int(v) for k, v in lookback_period.items()}
+    except Exception as e:
+        exception_msg = "Formatting lookback period failed"
+        raise ValueError(exception_msg) from e
+
+
 class GithubAuditLogExtractor(Extractor):
     """
     Extracts audit logs from the GitHub REST API.
@@ -46,12 +119,17 @@ def __init__(
         self.exclude_actors = exclude_actors
 
     async def extract_records(self) -> AsyncGenerator[GithubAuditLog]:
-        async for audit in self.client.fetch_enterprise_audit_log(
-            self.enterprise_name,
-            self.actions,
-            self.actors,
-            self.exclude_actors,
-            self.lookback_period,
-        ):
-            audit["timestamp"] = audit.pop("@timestamp")
-            yield audit
+        dates = generate_date_range(self.lookback_period) or [None]
+        for target_date in dates:
+            search_phrase = build_search_phrase(
+                actions=self.actions,
+                actors=self.actors,
+                exclude_actors=self.exclude_actors,
+                target_date=target_date,
+            )
+            async for audit in self.client.fetch_enterprise_audit_log(
+                self.enterprise_name,
+                search_phrase,
+            ):
+                audit["timestamp"] = audit.pop("@timestamp")
+                yield audit

nodestream_github/client/githubclient.py

@@ -6,12 +6,10 @@
 import json
 import logging
 from collections.abc import AsyncGenerator
-from datetime import UTC, datetime, timedelta
 from enum import Enum
 from typing import Any
 
 import httpx
-from dateutil.relativedelta import relativedelta
 from limits import RateLimitItem, RateLimitItemPerMinute
 from limits.aio.storage import MemoryStorage
 from limits.aio.strategies import MovingWindowRateLimiter, RateLimiter
@@ -74,77 +72,6 @@ def _fetch_problem(title: str, e: httpx.HTTPError):
             logger.warning("Problem fetching %s", title, exc_info=e, stacklevel=2)
 
 
-def validate_lookback_period(lookback_period: dict[str, int]) -> dict[str, int]:
-    """Sanitize the lookback period to only include valid keys."""
-
-    def validate_positive_int(value: int) -> int:
-        converted = int(value)
-        if converted <= 0:
-            negative_value_exception_msg = (
-                f"Lookback period values must be positive: {value}"
-            )
-            raise ValueError(negative_value_exception_msg)
-        return converted
-
-    try:
-        return {k: validate_positive_int(v) for k, v in lookback_period.items()}
-    except Exception as e:
-        exception_msg = "Formatting lookback period failed"
-        raise ValueError(exception_msg) from e
-
-
-def generate_date_range(lookback_period: dict[str, int]) -> list[str]:
-    """
-    Generate a list of date strings in YYYY-MM-DD format for
-    the given lookback period.
-    """
-    if not lookback_period:
-        return []
-
-    end_date = datetime.now(tz=UTC).date()
-    start_date = (datetime.now(tz=UTC) - relativedelta(**lookback_period)).date()
-
-    delta_days = (end_date - start_date).days + 1
-    return [
-        (start_date + timedelta(days=i)).strftime("%Y-%m-%d") for i in range(delta_days)
-    ]
-
-
-def build_search_phrase(
-    actions: list[str],
-    actors: list[str],
-    exclude_actors: list[str],
-    target_date: str | None = None,
-) -> str:
-    # adding action-based filtering
-    actions_phrase = ""
-    if actions:
-        actions_phrase = " ".join(f"action:{action}" for action in actions)
-
-    # adding date-based filtering for a specific date
-    date_filter = f"created:{target_date}" if target_date else ""
-
-    # adding actor-based filtering
-    actors_phrase = ""
-    if actors:
-        actors_phrase = " ".join(f"actor:{actor}" for actor in actors)
-
-    # adding exclude_actors based filtering
-    exclude_actors_phrase = ""
-    if exclude_actors:
-        exclude_actors_phrase = " ".join(f"-actor:{actor}" for actor in exclude_actors)
-    return " ".join(
-        section
-        for section in [
-            actions_phrase,
-            date_filter,
-            actors_phrase,
-            exclude_actors_phrase,
-        ]
-        if section
-    ).strip()
-
-
 class GithubRestApiClient:
     def __init__(
         self,
@@ -411,30 +338,17 @@ async def fetch_all_organizations(self) -> AsyncGenerator[types.GithubOrg]:
     async def fetch_enterprise_audit_log(
         self,
         enterprise_name: str,
-        actions: list[str],
-        actors: list[str],
-        exclude_actors: list[str],
-        lookback_period: dict[str, int],
+        search_phrase: str | None = None,
     ) -> AsyncGenerator[types.GithubAuditLog]:
         """Fetches enterprise-wide audit log data
-
-        https://docs.github.com/en/enterprise-cloud@latest/rest/enterprise-admin/audit-log?apiVersion=2022-11-28#get-the-audit-log-for-an-enterprise
+        https://docs.github.com/en/[email protected]/rest/enterprise-admin/audit-log?apiVersion=2022-11-28#get-the-audit-log-for-an-enterprise
         """
         try:
-            dates = generate_date_range(lookback_period) or [None]
-
-            for target_date in dates:
-                search_phrase = build_search_phrase(
-                    actions=actions,
-                    actors=actors,
-                    exclude_actors=exclude_actors,
-                    target_date=target_date,
-                )
-                params = {"phrase": search_phrase} if search_phrase else {}
-                async for audit in self._get_paginated(
-                    f"enterprises/{enterprise_name}/audit-log", params=params
-                ):
-                    yield audit
+            params = {"phrase": search_phrase} if search_phrase else {}
+            async for audit in self._get_paginated(
+                f"enterprises/{enterprise_name}/audit-log", params=params
+            ):
+                yield audit
         except httpx.HTTPError as e:
             _fetch_problem("audit log", e)
 

tests/client/test_githubclient.py

@@ -6,8 +6,6 @@
 from nodestream_github.client.githubclient import (
     GithubRestApiClient,
     RateLimitedError,
-    generate_date_range,
-    validate_lookback_period,
 )
 from tests.mocks.githubrest import DEFAULT_BASE_URL, DEFAULT_HOSTNAME
 
@@ -133,124 +131,3 @@ async def test_pagination_truncate_warning(
 def test_all_null_args():
     # noinspection PyTypeChecker
     assert GithubRestApiClient(auth_token=None, github_hostname=None)
-
-
-# Test generate_date_range
-
-
-def test_generate_date_range_empty_lookback_period():
-    result = generate_date_range({})
-    assert result == []
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_days_only():
-    result = generate_date_range({"days": 3})
-    expected = ["2025-07-29", "2025-07-30", "2025-07-31", "2025-08-01"]
-    assert result == expected
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_zero_days():
-    """Test with zero days (same day only)."""
-    result = generate_date_range({"days": 0})
-    expected = ["2025-08-01"]
-    assert result == expected
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_months_only():
-    result = generate_date_range({"months": 1})
-    # July 1 to August 1 = 32 days
-    assert len(result) == 32
-    assert result[0] == "2025-07-01"
-    assert result[-1] == "2025-08-01"
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_years_only():
-    result = generate_date_range({"years": 1})
-    assert len(result) == 366
-    assert result[0] == "2024-08-01"
-    assert result[-1] == "2025-08-01"
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_combined_periods():
-    """Test with combined periods."""
-    result = generate_date_range({"months": 1, "days": 5})
-    assert result[0] == "2025-06-26"
-    assert result[-1] == "2025-08-01"
-
-
-@freeze_time("2025-08-01")
-def test_generate_date_range_complex_combination():
-    """Test with years, months, and days."""
-    result = generate_date_range({"years": 1, "months": 2, "days": 10})
-    assert result[0] == "2024-05-22"
-    assert result[-1] == "2025-08-01"
-
-
-# Test validate_lookback_period
-
-
-def test_validate_lookback_period_valid_input():
-    result = validate_lookback_period({"days": 7, "months": 2, "years": 1})
-    expected = {"days": 7, "months": 2, "years": 1}
-    assert result == expected
-
-
-def test_validate_lookback_period_empty_dict():
-    result = validate_lookback_period({})
-    assert result == {}
-
-
-def test_validate_lookback_period_single_value():
-    result = validate_lookback_period({"days": 30})
-    assert result == {"days": 30}
-
-
-def test_validate_lookback_period_string_to_int_conversion():
-    result = validate_lookback_period({"days": "7", "months": "2"})
-    expected = {"days": 7, "months": 2}
-    assert result == expected
-
-
-def test_validate_lookback_period_zero_value():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": 0})
-
-
-def test_validate_lookback_period_negative_value():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": -5})
-
-
-def test_validate_lookback_period_multiple_negative_values():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": 7, "months": -1, "years": 2})
-
-
-def test_validate_lookback_period_zero_string():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": "0"})
-
-
-def test_validate_lookback_period_negative_string():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"months": "-10"})
-
-
-def test_validate_lookback_period_invalid_string():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": "invalid"})
-
-
-def test_validate_lookback_period_invalid_type():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": []})
-
-
-def test_validate_lookback_period_none_value():
-    with pytest.raises(ValueError, match="Formatting lookback period failed"):
-        validate_lookback_period({"days": None})