Skip to content

Commit

Permalink
refactor: Use tz-aware datetimes (#332)
Browse files Browse the repository at this point in the history
  • Loading branch information
@edgarrmondragon
edgarrmondragon authored Dec 23, 2024
1 parent ad5db0a commit 29d396f
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 27 deletions.
22 changes: 11 additions & 11 deletions tap_github/authenticator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import logging
import time
from copy import deepcopy
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
from os import environ
from random import choice, shuffle
from typing import Any
Expand Down Expand Up @@ -49,8 +49,9 @@ def __init__(
def update_rate_limit(self, response_headers: Any) -> None:
self.rate_limit = int(response_headers["X-RateLimit-Limit"])
self.rate_limit_remaining = int(response_headers["X-RateLimit-Remaining"])
self.rate_limit_reset = datetime.utcfromtimestamp(
int(response_headers["X-RateLimit-Reset"])
self.rate_limit_reset = datetime.fromtimestamp(
int(response_headers["X-RateLimit-Reset"]),
tz=timezone.utc,
)
self.rate_limit_used = int(response_headers["X-RateLimit-Used"])

Expand Down Expand Up @@ -86,10 +87,9 @@ def has_calls_remaining(self) -> bool:
"""
if self.rate_limit_reset is None:
return True
return (
self.rate_limit_used <= (self.rate_limit - self.rate_limit_buffer)
or self.rate_limit_reset <= datetime.now()
)
return self.rate_limit_used <= (
self.rate_limit - self.rate_limit_buffer
) or self.rate_limit_reset <= datetime.now(tz=timezone.utc)


class PersonalTokenManager(TokenManager):
Expand Down Expand Up @@ -126,7 +126,7 @@ def generate_app_access_token(
github_private_key: str,
github_installation_id: str | None = None,
) -> tuple[str, datetime]:
produced_at = datetime.now()
produced_at = datetime.now(tz=timezone.utc)
jwt_token = generate_jwt_token(github_app_id, github_private_key)

headers = {"Authorization": f"Bearer {jwt_token}"}
Expand Down Expand Up @@ -219,9 +219,9 @@ def has_calls_remaining(self) -> bool:
True if the token is valid and has enough api calls remaining.
"""
if self.token_expires_at is not None:
close_to_expiry = datetime.now() > self.token_expires_at - timedelta(
minutes=self.expiry_time_buffer
)
close_to_expiry = datetime.now(
tz=timezone.utc
) > self.token_expires_at - timedelta(minutes=self.expiry_time_buffer)

if close_to_expiry:
self.claim_token()
Expand Down
48 changes: 32 additions & 16 deletions tap_github/tests/test_authenticator.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import re
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
from unittest.mock import MagicMock, call, patch

import pytest
Expand All @@ -14,6 +14,10 @@
)


def _now():
return datetime.now(tz=timezone.utc)


class TestTokenManager:
def test_default_rate_limits(self):
token_manager = TokenManager("mytoken", rate_limit_buffer=700)
Expand All @@ -40,7 +44,15 @@ def test_update_rate_limit(self):

assert token_manager.rate_limit == 5000
assert token_manager.rate_limit_remaining == 4999
assert token_manager.rate_limit_reset == datetime(2013, 7, 1, 17, 47, 53)
assert token_manager.rate_limit_reset == datetime(
2013,
7,
1,
17,
47,
53,
tzinfo=timezone.utc,
)
assert token_manager.rate_limit_used == 1

def test_is_valid_token_successful(self):
Expand Down Expand Up @@ -108,9 +120,7 @@ def test_has_calls_remaining_fails_if_few_calls_remaining_and_reset_time_not_rea
mock_response_headers = {
"X-RateLimit-Limit": "5000",
"X-RateLimit-Remaining": "1",
"X-RateLimit-Reset": str(
int((datetime.now() + timedelta(days=100)).timestamp())
),
"X-RateLimit-Reset": str(int((_now() + timedelta(days=100)).timestamp())),
"X-RateLimit-Used": "4999",
}

Expand Down Expand Up @@ -164,8 +174,8 @@ def test_successful_token_generation(self):
assert token_manager.token_expires_at == token_time

def test_has_calls_remaining_regenerates_a_token_if_close_to_expiry(self):
unexpired_time = datetime.now() + timedelta(days=1)
expired_time = datetime.now() - timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
expired_time = _now() - timedelta(days=1)
with patch.object(AppTokenManager, "is_valid_token", return_value=True), patch(
"tap_github.authenticator.generate_app_access_token",
return_value=("valid_token", unexpired_time),
Expand Down Expand Up @@ -193,8 +203,8 @@ def test_has_calls_remaining_regenerates_a_token_if_close_to_expiry(self):
)

def test_has_calls_remaining_logs_warning_if_token_regeneration_fails(self):
unexpired_time = datetime.now() + timedelta(days=1)
expired_time = datetime.now() - timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
expired_time = _now() - timedelta(days=1)
with patch.object(
AppTokenManager, "is_valid_token", return_value=True
) as mock_is_valid, patch(
Expand Down Expand Up @@ -222,7 +232,7 @@ def test_has_calls_remaining_logs_warning_if_token_regeneration_fails(self):
)

def test_has_calls_remaining_succeeds_if_token_new_and_never_used(self):
unexpired_time = datetime.now() + timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
with patch.object(AppTokenManager, "is_valid_token", return_value=True), patch(
"tap_github.authenticator.generate_app_access_token",
return_value=("valid_token", unexpired_time),
Expand All @@ -231,7 +241,7 @@ def test_has_calls_remaining_succeeds_if_token_new_and_never_used(self):
assert token_manager.has_calls_remaining()

def test_has_calls_remaining_succeeds_if_time_and_requests_left(self):
unexpired_time = datetime.now() + timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
with patch.object(AppTokenManager, "is_valid_token", return_value=True), patch(
"tap_github.authenticator.generate_app_access_token",
return_value=("valid_token", unexpired_time),
Expand All @@ -249,7 +259,7 @@ def test_has_calls_remaining_succeeds_if_time_and_requests_left(self):
assert token_manager.has_calls_remaining()

def test_has_calls_remaining_succeeds_if_time_left_and_reset_time_reached(self):
unexpired_time = datetime.now() + timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
with patch.object(AppTokenManager, "is_valid_token", return_value=True), patch(
"tap_github.authenticator.generate_app_access_token",
return_value=("valid_token", unexpired_time),
Expand All @@ -271,7 +281,7 @@ def test_has_calls_remaining_succeeds_if_time_left_and_reset_time_reached(self):
def test_has_calls_remaining_fails_if_time_left_and_few_calls_remaining_and_reset_time_not_reached( # noqa: E501
self,
):
unexpired_time = datetime.now() + timedelta(days=1)
unexpired_time = _now() + timedelta(days=1)
with patch.object(AppTokenManager, "is_valid_token", return_value=True), patch(
"tap_github.authenticator.generate_app_access_token",
return_value=("valid_token", unexpired_time),
Expand All @@ -280,7 +290,7 @@ def test_has_calls_remaining_fails_if_time_left_and_few_calls_remaining_and_rese
"X-RateLimit-Limit": "5000",
"X-RateLimit-Remaining": "1",
"X-RateLimit-Reset": str(
int((datetime.now() + timedelta(days=100)).timestamp())
int((_now() + timedelta(days=100)).timestamp())
),
"X-RateLimit-Used": "4999",
}
Expand Down Expand Up @@ -427,7 +437,10 @@ def test_all_token_types(self, mock_stream):
):
stream = mock_stream
stream.config.update(
{"auth_token": "gt5", "additional_auth_tokens": ["gt7", "gt8", "gt9"]}
{
"auth_token": "gt5",
"additional_auth_tokens": ["gt7", "gt8", "gt9"],
}
)
auth = GitHubTokenAuthenticator(stream=stream)
token_managers = auth.prepare_tokens()
Expand Down Expand Up @@ -568,7 +581,10 @@ def test_prepare_tokens_returns_empty_if_all_tokens_invalid(self, mock_stream):
):
stream = mock_stream
stream.config.update(
{"auth_token": "gt5", "additional_auth_tokens": ["gt7", "gt8", "gt9"]}
{
"auth_token": "gt5",
"additional_auth_tokens": ["gt7", "gt8", "gt9"],
}
)
auth = GitHubTokenAuthenticator(stream=stream)
token_managers = auth.prepare_tokens()
Expand Down

0 comments on commit 29d396f

Please sign in to comment.