refactor: Use tz-aware datetimes (#332)

Author: edgarrmondragon

tap_github/authenticator.py

@@ -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
@@ -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"])
 
@@ -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):
@@ -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}"}
@@ -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()

tap_github/tests/test_authenticator.py

@@ -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
@@ -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)
@@ -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):
@@ -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",
         }
 
@@ -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),
@@ -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(
@@ -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),
@@ -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),
@@ -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),
@@ -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),
@@ -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",
             }
@@ -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()
@@ -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()