From cf856d07adfd4a56fdaeb4ab09b39b398c939c7c Mon Sep 17 00:00:00 2001 From: Mariatta Wijaya Date: Wed, 10 May 2023 17:01:08 -0700 Subject: [PATCH 1/4] feat: Replace utcnow and utcfromtimestamp These will be deprecated starting in Python 3.12. Replaced: - `datetime.datetime.utcnow()` with `datetime.datetime.now(tz=datetime.timezone.utc)` - `datetime.utcfromtimestamp()` with `datetime.fromtimestamp(tz=datetime.timezone.utc)` --- google/cloud/firestore_v1/_helpers.py | 4 +++- google/cloud/firestore_v1/bulk_writer.py | 8 ++++++-- google/cloud/firestore_v1/rate_limiter.py | 2 +- tests/system/test_system.py | 9 ++++----- tests/system/test_system_async.py | 5 ++--- tests/unit/v1/test__helpers.py | 13 ++++++++----- tests/unit/v1/test_async_client.py | 2 +- tests/unit/v1/test_base_client.py | 2 +- tests/unit/v1/test_base_query.py | 2 +- tests/unit/v1/test_client.py | 2 +- tests/unit/v1/test_rate_limiter.py | 2 +- 11 files changed, 29 insertions(+), 22 deletions(-) diff --git a/google/cloud/firestore_v1/_helpers.py b/google/cloud/firestore_v1/_helpers.py index 3b6b7886bc..d34dfc5735 100644 --- a/google/cloud/firestore_v1/_helpers.py +++ b/google/cloud/firestore_v1/_helpers.py @@ -1128,7 +1128,9 @@ def build_timestamp( dt: Optional[Union[DatetimeWithNanoseconds, datetime.datetime]] = None ) -> Timestamp: """Returns the supplied datetime (or "now") as a Timestamp""" - return _datetime_to_pb_timestamp(dt or DatetimeWithNanoseconds.utcnow()) + return _datetime_to_pb_timestamp( + dt or DatetimeWithNanoseconds.now(tz=datetime.timezone.utc) + ) def compare_timestamps( diff --git a/google/cloud/firestore_v1/bulk_writer.py b/google/cloud/firestore_v1/bulk_writer.py index 9c7c0d5c9e..b736d88781 100644 --- a/google/cloud/firestore_v1/bulk_writer.py +++ b/google/cloud/firestore_v1/bulk_writer.py @@ -192,7 +192,9 @@ def _retry_operation( elif self._options.retry == BulkRetry.linear: delay = operation.attempts - run_at = datetime.datetime.utcnow() + datetime.timedelta(seconds=delay) + run_at = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta( + seconds=delay + ) # Use of `bisect.insort` maintains the requirement that `self._retries` # always remain sorted by each object's `run_at` time. Note that it is @@ -498,7 +500,9 @@ def _schedule_ready_retries(self): # ever adding to it via `bisect.insort`), and because `OperationRetry` # objects are comparable against `datetime` objects, this bisect functionally # returns the number of retires that are ready for immediate reenlistment. - take_until_index = bisect.bisect(self._retries, datetime.datetime.utcnow()) + take_until_index = bisect.bisect( + self._retries, datetime.datetime.now(tz=datetime.timezone.utc) + ) for _ in range(take_until_index): retry: OperationRetry = self._retries.popleft() diff --git a/google/cloud/firestore_v1/rate_limiter.py b/google/cloud/firestore_v1/rate_limiter.py index 2543869532..ebf728ef86 100644 --- a/google/cloud/firestore_v1/rate_limiter.py +++ b/google/cloud/firestore_v1/rate_limiter.py @@ -17,7 +17,7 @@ def utcnow(): - return datetime.datetime.utcnow() + return datetime.datetime.now(tz=datetime.timezone.utc) default_initial_tokens: int = 500 diff --git a/tests/system/test_system.py b/tests/system/test_system.py index eac329bcb3..1dde681d92 100644 --- a/tests/system/test_system.py +++ b/tests/system/test_system.py @@ -26,7 +26,6 @@ from google.api_core.exceptions import InvalidArgument from google.api_core.exceptions import NotFound from google.cloud._helpers import _datetime_to_pb_timestamp -from google.cloud._helpers import UTC from google.cloud import firestore_v1 as firestore from google.cloud.firestore_v1.base_query import FieldFilter, And, Or @@ -90,7 +89,7 @@ def test_collections_w_import(): def test_create_document(client, cleanup): - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) collection_id = "doc-create" + UNIQUE_RESOURCE_ID document_id = "doc" + UNIQUE_RESOURCE_ID document = client.document(collection_id, document_id) @@ -364,7 +363,7 @@ def check_snapshot(snapshot, document, data, write_result): def test_document_get(client, cleanup): - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) document_id = "for-get" + UNIQUE_RESOURCE_ID document = client.document("created", document_id) # Add to clean-up before API request (in case ``create()`` fails). @@ -1637,7 +1636,7 @@ def on_snapshot(docs, changes, read_time): def test_repro_429(client, cleanup): # See: https://github.com/googleapis/python-firestore/issues/429 - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) collection = client.collection("repro-429" + UNIQUE_RESOURCE_ID) for document_id in [f"doc-{doc_id:02d}" for doc_id in range(30)]: @@ -1664,7 +1663,7 @@ def test_repro_429(client, cleanup): def test_repro_391(client, cleanup): # See: https://github.com/googleapis/python-firestore/issues/391 - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) collection = client.collection("repro-391" + UNIQUE_RESOURCE_ID) document_ids = [f"doc-{doc_id:02d}" for doc_id in range(30)] diff --git a/tests/system/test_system_async.py b/tests/system/test_system_async.py index 9b25039fc3..74b8b1749d 100644 --- a/tests/system/test_system_async.py +++ b/tests/system/test_system_async.py @@ -33,7 +33,6 @@ from google.api_core.exceptions import InvalidArgument from google.api_core.exceptions import NotFound from google.cloud._helpers import _datetime_to_pb_timestamp -from google.cloud._helpers import UTC from google.cloud import firestore_v1 as firestore from google.cloud.firestore_v1.base_query import FieldFilter, And, Or @@ -118,7 +117,7 @@ async def test_collections_w_import(): async def test_create_document(client, cleanup): - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) collection_id = "doc-create" + UNIQUE_RESOURCE_ID document_id = "doc" + UNIQUE_RESOURCE_ID document = client.document(collection_id, document_id) @@ -394,7 +393,7 @@ def check_snapshot(snapshot, document, data, write_result): async def test_document_get(client, cleanup): - now = datetime.datetime.utcnow().replace(tzinfo=UTC) + now = datetime.datetime.now(tz=datetime.timezone.utc) document_id = "for-get" + UNIQUE_RESOURCE_ID document = client.document("created", document_id) # Add to clean-up before API request (in case ``create()`` fails). diff --git a/tests/unit/v1/test__helpers.py b/tests/unit/v1/test__helpers.py index 0a6dee40e3..c808625b92 100644 --- a/tests/unit/v1/test__helpers.py +++ b/tests/unit/v1/test__helpers.py @@ -205,7 +205,9 @@ def test_encode_value_w_datetime_wo_nanos(): dt_nanos = 458816000 # Make sure precision is valid in microseconds too. assert dt_nanos % 1000 == 0 - dt_val = datetime.datetime.utcfromtimestamp(dt_seconds + 1e-9 * dt_nanos) + dt_val = datetime.datetime.fromtimestamp( + dt_seconds + 1e-9 * dt_nanos, tz=datetime.timezone.utc + ) result = encode_value(dt_val) timestamp_pb = timestamp_pb2.Timestamp(seconds=dt_seconds, nanos=dt_nanos) @@ -304,7 +306,9 @@ def test_encode_dict_w_many_types(): dt_nanos = 465964000 # Make sure precision is valid in microseconds too. assert dt_nanos % 1000 == 0 - dt_val = datetime.datetime.utcfromtimestamp(dt_seconds + 1e-9 * dt_nanos) + dt_val = datetime.datetime.fromtimestamp( + dt_seconds + 1e-9 * dt_nanos, tz=datetime.timezone.utc + ) client = _make_client() document = client.document("most", "adjective", "thing", "here") @@ -646,7 +650,6 @@ def test_decode_dict_w_many_types(): from google.protobuf import timestamp_pb2 from google.cloud.firestore_v1.types.document import ArrayValue from google.cloud.firestore_v1.types.document import MapValue - from google.cloud._helpers import UTC from google.cloud.firestore_v1.field_path import FieldPath from google.cloud.firestore_v1._helpers import decode_dict @@ -654,8 +657,8 @@ def test_decode_dict_w_many_types(): dt_nanos = 667285000 # Make sure precision is valid in microseconds too. assert dt_nanos % 1000 == 0 - dt_val = datetime.datetime.utcfromtimestamp(dt_seconds + 1e-9 * dt_nanos).replace( - tzinfo=UTC + dt_val = datetime.datetime.fromtimestamp( + dt_seconds + 1e-9 * dt_nanos, tz=datetime.timezone.utc ) value_fields = { diff --git a/tests/unit/v1/test_async_client.py b/tests/unit/v1/test_async_client.py index 69785f5b82..f3faf40dbf 100644 --- a/tests/unit/v1/test_async_client.py +++ b/tests/unit/v1/test_async_client.py @@ -554,7 +554,7 @@ def _doc_get_info(ref_string, values): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers - now = datetime.datetime.utcnow() + now = datetime.datetime.now(tz=datetime.timezone.utc) read_time = _datetime_to_pb_timestamp(now) delta = datetime.timedelta(seconds=100) update_time = _datetime_to_pb_timestamp(now - delta) diff --git a/tests/unit/v1/test_base_client.py b/tests/unit/v1/test_base_client.py index dfc235641d..57d278daa2 100644 --- a/tests/unit/v1/test_base_client.py +++ b/tests/unit/v1/test_base_client.py @@ -404,7 +404,7 @@ def test__parse_batch_get_found(): from google.cloud.firestore_v1.document import DocumentSnapshot from google.cloud.firestore_v1.base_client import _parse_batch_get - now = datetime.datetime.utcnow() + now = datetime.datetime.now(tz=datetime.timezone.utc) read_time = _datetime_to_pb_timestamp(now) delta = datetime.timedelta(seconds=100) update_time = _datetime_to_pb_timestamp(now - delta) diff --git a/tests/unit/v1/test_base_query.py b/tests/unit/v1/test_base_query.py index 4b8093f1a7..4b1ea79c25 100644 --- a/tests/unit/v1/test_base_query.py +++ b/tests/unit/v1/test_base_query.py @@ -1941,7 +1941,7 @@ def _make_query_response(**kwargs): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers - now = datetime.datetime.utcnow() + now = datetime.datetime.now(tz=datetime.timezone.utc) read_time = _datetime_to_pb_timestamp(now) kwargs["read_time"] = read_time diff --git a/tests/unit/v1/test_client.py b/tests/unit/v1/test_client.py index 563419b30d..7fb09c7dd0 100644 --- a/tests/unit/v1/test_client.py +++ b/tests/unit/v1/test_client.py @@ -552,7 +552,7 @@ def _doc_get_info(ref_string, values): from google.cloud._helpers import _datetime_to_pb_timestamp from google.cloud.firestore_v1 import _helpers - now = datetime.datetime.utcnow() + now = datetime.datetime.now(tz=datetime.timezone.utc) read_time = _datetime_to_pb_timestamp(now) delta = datetime.timedelta(seconds=100) update_time = _datetime_to_pb_timestamp(now - delta) diff --git a/tests/unit/v1/test_rate_limiter.py b/tests/unit/v1/test_rate_limiter.py index d27b7ee810..b4cdd8fb7c 100644 --- a/tests/unit/v1/test_rate_limiter.py +++ b/tests/unit/v1/test_rate_limiter.py @@ -19,7 +19,7 @@ # Pick a point in time as the center of our universe for this test run. # It is okay for this to update every time the tests are run. -fake_now = datetime.datetime.utcnow() +fake_now = datetime.datetime.now(tz=datetime.timezone.utc) def now_plus_n(seconds: int = 0, microseconds: int = 0) -> datetime.timedelta: From 075e999a4638fd137e65757f8d5a65a3aa3fcef9 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 4 Dec 2023 16:32:25 +0000 Subject: [PATCH 2/4] docs: deprecate google.cloud.firestore_v1.rate_limiter.utcnow --- google/cloud/firestore_v1/rate_limiter.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/google/cloud/firestore_v1/rate_limiter.py b/google/cloud/firestore_v1/rate_limiter.py index ebf728ef86..7925887c88 100644 --- a/google/cloud/firestore_v1/rate_limiter.py +++ b/google/cloud/firestore_v1/rate_limiter.py @@ -14,10 +14,20 @@ import datetime from typing import NoReturn, Optional +import warnings def utcnow(): - return datetime.datetime.now(tz=datetime.timezone.utc) + """ + google.cloud.firestore_v1.rate_limiter.utcnow() is deprecated. + Use datetime.datetime.now(datetime.timezone.utc) instead. + """ + warnings.warn( + "google.cloud.firestore_v1.rate_limiter.utcnow() is deprecated. " + "Use datetime.datetime.now(datetime.timezone.utc) instead.", + DeprecationWarning, + ) + return datetime.datetime.utcnow() default_initial_tokens: int = 500 From 5f9f04f1eae89956cbcfe3c08e41638ac976e80f Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 4 Dec 2023 16:32:53 +0000 Subject: [PATCH 3/4] remove usage of google.cloud.firestore_v1.rate_limiter.utcnow in code --- google/cloud/firestore_v1/rate_limiter.py | 11 +- noxfile.py | 1 + owlbot.py | 2 +- tests/unit/v1/test_rate_limiter.py | 360 +++++++++++----------- 4 files changed, 194 insertions(+), 180 deletions(-) diff --git a/google/cloud/firestore_v1/rate_limiter.py b/google/cloud/firestore_v1/rate_limiter.py index 7925887c88..8ca98dbe88 100644 --- a/google/cloud/firestore_v1/rate_limiter.py +++ b/google/cloud/firestore_v1/rate_limiter.py @@ -106,8 +106,9 @@ def __init__( self._phase: int = 0 def _start_clock(self): - self._start = self._start or utcnow() - self._last_refill = self._last_refill or utcnow() + utcnow = datetime.datetime.now(datetime.timezone.utc) + self._start = self._start or utcnow + self._last_refill = self._last_refill or utcnow def take_tokens(self, num: Optional[int] = 1, allow_less: bool = False) -> int: """Returns the number of available tokens, up to the amount requested.""" @@ -133,7 +134,9 @@ def _check_phase(self): This is a no-op unless a new [_phase_length] number of seconds since the start was crossed since it was last called. """ - age: datetime.timedelta = utcnow() - self._start + age: datetime.timedelta = ( + datetime.datetime.now(datetime.timezone.utc) - self._start + ) # Uses integer division to calculate the expected phase. We start in # Phase 0, so until [_phase_length] seconds have passed, this will @@ -162,7 +165,7 @@ def _increase_maximum_tokens(self) -> NoReturn: def _refill(self) -> NoReturn: """Replenishes any tokens that should have regenerated since the last operation.""" - now: datetime.datetime = utcnow() + now: datetime.datetime = datetime.datetime.now(datetime.timezone.utc) time_since_last_refill: datetime.timedelta = now - self._last_refill if time_since_last_refill: diff --git a/noxfile.py b/noxfile.py index 7bbf746683..b4b7202578 100644 --- a/noxfile.py +++ b/noxfile.py @@ -45,6 +45,7 @@ UNIT_TEST_EXTERNAL_DEPENDENCIES: List[str] = [ "aiounittest", "six", + "freezegun", ] UNIT_TEST_LOCAL_DEPENDENCIES: List[str] = [] UNIT_TEST_DEPENDENCIES: List[str] = [] diff --git a/owlbot.py b/owlbot.py index 812eadab06..4384bb53a6 100644 --- a/owlbot.py +++ b/owlbot.py @@ -139,7 +139,7 @@ def update_fixup_scripts(library): templated_files = common.py_library( samples=False, # set to True only if there are samples system_test_python_versions=["3.7"], - unit_test_external_dependencies=["aiounittest", "six"], + unit_test_external_dependencies=["aiounittest", "six", "freezegun"], system_test_external_dependencies=["pytest-asyncio", "six"], microgenerator=True, cov_level=100, diff --git a/tests/unit/v1/test_rate_limiter.py b/tests/unit/v1/test_rate_limiter.py index b4cdd8fb7c..a41519b4f0 100644 --- a/tests/unit/v1/test_rate_limiter.py +++ b/tests/unit/v1/test_rate_limiter.py @@ -14,207 +14,217 @@ import datetime -import mock +import freezegun +from google.cloud.firestore_v1 import rate_limiter # Pick a point in time as the center of our universe for this test run. # It is okay for this to update every time the tests are run. fake_now = datetime.datetime.now(tz=datetime.timezone.utc) -def now_plus_n(seconds: int = 0, microseconds: int = 0) -> datetime.timedelta: - return fake_now + datetime.timedelta( - seconds=seconds, - microseconds=microseconds, - ) - - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_rate_limiter_basic(mocked_now): +def test_rate_limiter_basic(): """Verifies that if the clock does not advance, the RateLimiter allows 500 writes before crashing out. """ - from google.cloud.firestore_v1 import rate_limiter + with freezegun.freeze_time(fake_now): + # This RateLimiter will never advance. + ramp = rate_limiter.RateLimiter() + for _ in range(rate_limiter.default_initial_tokens): + assert ramp.take_tokens() == 1 + assert ramp.take_tokens() == 0 - mocked_now.return_value = fake_now - # This RateLimiter will never advance. Poor fella. - ramp = rate_limiter.RateLimiter() - for _ in range(rate_limiter.default_initial_tokens): - assert ramp.take_tokens() == 1 - assert ramp.take_tokens() == 0 - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_rate_limiter_with_refill(mocked_now): +def test_rate_limiter_with_refill(): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ - from google.cloud.firestore_v1 import rate_limiter - - mocked_now.return_value = fake_now - ramp = rate_limiter.RateLimiter() - ramp._available_tokens = 0 - assert ramp.take_tokens() == 0 - # Advance the clock 0.1 seconds - mocked_now.return_value = now_plus_n(microseconds=100000) - for _ in range(round(rate_limiter.default_initial_tokens / 10)): - assert ramp.take_tokens() == 1 - assert ramp.take_tokens() == 0 - - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_rate_limiter_phase_length(mocked_now): + with freezegun.freeze_time(fake_now) as frozen_datetime: + ramp = rate_limiter.RateLimiter() + ramp._available_tokens = 0 + assert ramp.take_tokens() == 0 + # Advance the clock 0.1 seconds + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + microseconds=100000, + ) + ) + for _ in range(round(rate_limiter.default_initial_tokens / 10)): + assert ramp.take_tokens() == 1 + assert ramp.take_tokens() == 0 + + +def test_rate_limiter_phase_length(): """Verifies that if the clock advances, the RateLimiter allows appropriate additional writes. """ - from google.cloud.firestore_v1 import rate_limiter - - mocked_now.return_value = fake_now - ramp = rate_limiter.RateLimiter() - assert ramp.take_tokens() == 1 - ramp._available_tokens = 0 - assert ramp.take_tokens() == 0 - # Advance the clock 1 phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length, - microseconds=1, - ) - for _ in range(round(rate_limiter.default_initial_tokens * 3 / 2)): - assert ramp.take_tokens() - - assert ramp.take_tokens() == 0 - - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_rate_limiter_idle_phase_length(mocked_now): - """Verifies that if the clock advances but nothing happens, the RateLimiter - doesn't ramp up. - """ - from google.cloud.firestore_v1 import rate_limiter - - mocked_now.return_value = fake_now - ramp = rate_limiter.RateLimiter() - ramp._available_tokens = 0 - assert ramp.take_tokens() == 0 - # Advance the clock 1 phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length, - microseconds=1, - ) - for _ in range(round(rate_limiter.default_initial_tokens)): + with freezegun.freeze_time(fake_now) as frozen_datetime: + ramp = rate_limiter.RateLimiter() assert ramp.take_tokens() == 1 - assert ramp._maximum_tokens == 500 - assert ramp.take_tokens() == 0 + ramp._available_tokens = 0 + assert ramp.take_tokens() == 0 + + # Advance the clock 1 phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length, + microseconds=1, + ) + ) + for _ in range(round(rate_limiter.default_initial_tokens * 3 / 2)): + assert ramp.take_tokens() + assert ramp.take_tokens() == 0 -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_take_batch_size(mocked_now): + +def test_rate_limiter_idle_phase_length(): """Verifies that if the clock advances but nothing happens, the RateLimiter doesn't ramp up. """ - from google.cloud.firestore_v1 import rate_limiter - - page_size: int = 20 - mocked_now.return_value = fake_now - ramp = rate_limiter.RateLimiter() - ramp._available_tokens = 15 - assert ramp.take_tokens(page_size, allow_less=True) == 15 - # Advance the clock 1 phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length, - microseconds=1, - ) - ramp._check_phase() - assert ramp._maximum_tokens == 750 - - for _ in range(740 // page_size): - assert ramp.take_tokens(page_size) == page_size - assert ramp.take_tokens(page_size, allow_less=True) == 10 - assert ramp.take_tokens(page_size, allow_less=True) == 0 - - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_phase_progress(mocked_now): - from google.cloud.firestore_v1 import rate_limiter - - mocked_now.return_value = fake_now - - ramp = rate_limiter.RateLimiter() - assert ramp._phase == 0 - assert ramp._maximum_tokens == 500 - ramp.take_tokens() - - # Advance the clock 1 phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length, - microseconds=1, - ) - ramp.take_tokens() - assert ramp._phase == 1 - assert ramp._maximum_tokens == 750 - - # Advance the clock another phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length * 2, - microseconds=1, - ) - ramp.take_tokens() - assert ramp._phase == 2 - assert ramp._maximum_tokens == 1125 - - # Advance the clock another ms and the phase should not advance - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length * 2, - microseconds=2, - ) - ramp.take_tokens() - assert ramp._phase == 2 - assert ramp._maximum_tokens == 1125 - - -@mock.patch("google.cloud.firestore_v1.rate_limiter.utcnow") -def test_global_max_tokens(mocked_now): - from google.cloud.firestore_v1 import rate_limiter - - mocked_now.return_value = fake_now - - ramp = rate_limiter.RateLimiter( - global_max_tokens=499, - ) - assert ramp._phase == 0 - assert ramp._maximum_tokens == 499 - ramp.take_tokens() - - # Advance the clock 1 phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length, - microseconds=1, - ) - ramp.take_tokens() - assert ramp._phase == 1 - assert ramp._maximum_tokens == 499 - - # Advance the clock another phase - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length * 2, - microseconds=1, - ) - ramp.take_tokens() - assert ramp._phase == 2 - assert ramp._maximum_tokens == 499 - - # Advance the clock another ms and the phase should not advance - mocked_now.return_value = now_plus_n( - seconds=rate_limiter.default_phase_length * 2, - microseconds=2, - ) - ramp.take_tokens() - assert ramp._phase == 2 - assert ramp._maximum_tokens == 499 + with freezegun.freeze_time(fake_now) as frozen_datetime: + ramp = rate_limiter.RateLimiter() + ramp._available_tokens = 0 + assert ramp.take_tokens() == 0 + + # Advance the clock 1 phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length, + microseconds=1, + ) + ) + + for _ in range(round(rate_limiter.default_initial_tokens)): + assert ramp.take_tokens() == 1 + assert ramp._maximum_tokens == 500 + assert ramp.take_tokens() == 0 + + +def test_take_batch_size(): + """Verifies that if the clock advances but nothing happens, the RateLimiter + doesn't ramp up. + """ + with freezegun.freeze_time(fake_now) as frozen_datetime: + page_size: int = 20 + + ramp = rate_limiter.RateLimiter() + ramp._available_tokens = 15 + assert ramp.take_tokens(page_size, allow_less=True) == 15 + + # Advance the clock 1 phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length, + microseconds=1, + ) + ) + ramp._check_phase() + assert ramp._maximum_tokens == 750 + + for _ in range(740 // page_size): + assert ramp.take_tokens(page_size) == page_size + assert ramp.take_tokens(page_size, allow_less=True) == 10 + assert ramp.take_tokens(page_size, allow_less=True) == 0 + + +def test_phase_progress(): + with freezegun.freeze_time(fake_now) as frozen_datetime: + ramp = rate_limiter.RateLimiter() + assert ramp._phase == 0 + assert ramp._maximum_tokens == 500 + ramp.take_tokens() + + # Advance the clock 1 phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length, + microseconds=1, + ) + ) + ramp.take_tokens() + assert ramp._phase == 1 + assert ramp._maximum_tokens == 750 + + # Advance the clock another phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length * 2, + microseconds=1, + ) + ) + + ramp.take_tokens() + assert ramp._phase == 2 + assert ramp._maximum_tokens == 1125 + + # Advance the clock another ms and the phase should not advance + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length * 2, + microseconds=2, + ) + ) + + ramp.take_tokens() + assert ramp._phase == 2 + assert ramp._maximum_tokens == 1125 + + +def test_global_max_tokens(): + with freezegun.freeze_time(fake_now) as frozen_datetime: + ramp = rate_limiter.RateLimiter( + global_max_tokens=499, + ) + assert ramp._phase == 0 + assert ramp._maximum_tokens == 499 + ramp.take_tokens() + + # Advance the clock 1 phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length, + microseconds=1, + ) + ) + ramp.take_tokens() + assert ramp._phase == 1 + assert ramp._maximum_tokens == 499 + + # Advance the clock another phase + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length * 2, + microseconds=1, + ) + ) + + ramp.take_tokens() + assert ramp._phase == 2 + assert ramp._maximum_tokens == 499 + + # Advance the clock another ms and the phase should not advance + frozen_datetime.move_to( + fake_now + + datetime.timedelta( + seconds=rate_limiter.default_phase_length * 2, + microseconds=2, + ) + ) + + ramp.take_tokens() + assert ramp._phase == 2 + assert ramp._maximum_tokens == 499 def test_utcnow(): - from google.cloud.firestore_v1 import rate_limiter - now = rate_limiter.utcnow() assert isinstance(now, datetime.datetime) From 0f93993f82cebc029a862b70ddc2bd3393a404a0 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Mon, 4 Dec 2023 16:46:28 +0000 Subject: [PATCH 4/4] filter deprecation warning for google.cloud.firestore_v1.rate_limiter.utcnow --- tests/unit/v1/test_rate_limiter.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/v1/test_rate_limiter.py b/tests/unit/v1/test_rate_limiter.py index a41519b4f0..c23b85ae03 100644 --- a/tests/unit/v1/test_rate_limiter.py +++ b/tests/unit/v1/test_rate_limiter.py @@ -13,6 +13,7 @@ # limitations under the License. import datetime +import pytest import freezegun @@ -226,5 +227,9 @@ def test_global_max_tokens(): def test_utcnow(): - now = rate_limiter.utcnow() + with pytest.warns( + DeprecationWarning, + match="google.cloud.firestore_v1.rate_limiter.utcnow", + ): + now = rate_limiter.utcnow() assert isinstance(now, datetime.datetime)