diff --git a/google/cloud/alloydb/connector/refresh_utils.py b/google/cloud/alloydb/connector/refresh_utils.py index 1e173c5c..c1214ae5 100644 --- a/google/cloud/alloydb/connector/refresh_utils.py +++ b/google/cloud/alloydb/connector/refresh_utils.py @@ -26,9 +26,7 @@ _refresh_buffer: int = 4 * 60 # 4 minutes -def _seconds_until_refresh( - expiration: datetime, now: datetime = datetime.now(timezone.utc) -) -> int: +def _seconds_until_refresh(expiration: datetime) -> int: """ Calculates the duration to wait before starting the next refresh. Usually the duration will be half of the time until certificate @@ -36,12 +34,11 @@ def _seconds_until_refresh( Args: expiration (datetime.datetime): Time of certificate expiration. - now (datetime.datetime): Current time (UTC) Returns: int: Time in seconds to wait before performing next refresh. """ - duration = int((expiration - now).total_seconds()) + duration = int((expiration - datetime.now(timezone.utc)).total_seconds()) # if certificate duration is less than 1 hour if duration < 3600: diff --git a/tests/unit/test_refresh_utils.py b/tests/unit/test_refresh_utils.py index 1462f77d..8543f468 100644 --- a/tests/unit/test_refresh_utils.py +++ b/tests/unit/test_refresh_utils.py @@ -16,6 +16,8 @@ from datetime import timedelta from datetime import timezone +import pytest + from google.cloud.alloydb.connector.refresh_utils import _seconds_until_refresh @@ -24,8 +26,14 @@ def test_seconds_until_refresh_over_1_hour() -> None: Test _seconds_until_refresh returns proper time in seconds. If expiration is over 1 hour, should return duration/2. """ - now = datetime.now() - assert _seconds_until_refresh(now + timedelta(minutes=62), now) == 31 * 60 + # using pytest.approx since sometimes can be off by a second + assert ( + pytest.approx( + _seconds_until_refresh(datetime.now(timezone.utc) + timedelta(minutes=62)), + 1, + ) + == 31 * 60 + ) def test_seconds_until_refresh_under_1_hour_over_4_mins() -> None: @@ -34,8 +42,14 @@ def test_seconds_until_refresh_under_1_hour_over_4_mins() -> None: If expiration is under 1 hour and over 4 minutes, should return duration-refresh_buffer (refresh_buffer = 4 minutes). """ - now = datetime.now(timezone.utc) - assert _seconds_until_refresh(now + timedelta(minutes=5), now) == 60 + # using pytest.approx since sometimes can be off by a second + assert ( + pytest.approx( + _seconds_until_refresh(datetime.now(timezone.utc) + timedelta(minutes=5)), + 1, + ) + == 60 + ) def test_seconds_until_refresh_under_4_mins() -> None: