Browse Source

ref: move from freezegun to time-machine (part 1: introduce) (#56216)

time-machine is faster since it does not iterate through modules to
patch / unpatch

it's also actively maintained and typed properly
anthony sottile 1 year ago
parent
commit
dc8c226a26

+ 1 - 0
requirements-dev-frozen.txt

@@ -185,6 +185,7 @@ statsd==3.3
 stripe==2.61.0
 structlog==22.1.0
 symbolic==12.2.0
+time-machine==2.13.0
 tokenize-rt==5.0.0
 tomli==2.0.1
 toronado==0.1.0

+ 1 - 0
requirements-dev.txt

@@ -3,6 +3,7 @@
 covdefaults>=2.3.0
 docker>=3.7.0,<3.8.0
 freezegun>=1.2.2
+time-machine>=2.13.0
 honcho>=1.1.0
 openapi-core>=0.14.2
 pytest>=7.2.1

+ 7 - 0
src/sentry/testutils/helpers/datetime.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import time
 from datetime import datetime, timedelta
 
+import time_machine
 from django.utils import timezone
 
 __all__ = ["iso_format", "before_now", "timestamp_format"]
@@ -28,3 +31,7 @@ class MockClock:
     def __call__(self):
         self.time += timedelta(seconds=1)
         return self.time
+
+
+def freeze_time(t: str | datetime = "2023-09-13 01:02:00") -> time_machine.travel:
+    return time_machine.travel(t, tick=False)

+ 0 - 3
src/sentry/testutils/pytest/sentry.py

@@ -10,7 +10,6 @@ from hashlib import md5
 from typing import TypeVar
 from unittest import mock
 
-import freezegun
 import pytest
 from django.conf import settings
 from sentry_sdk import Hub
@@ -273,8 +272,6 @@ def pytest_configure(config: pytest.Config) -> None:
     # force celery registration
     from sentry.celery import app  # NOQA
 
-    freezegun.configure(extend_ignore_list=["sentry.utils.retries"])  # type: ignore[attr-defined]
-
 
 def register_extensions() -> None:
     from sentry.plugins.base import plugins

+ 6 - 6
tests/sentry/dynamic_sampling/tasks/test_common.py

@@ -2,7 +2,6 @@ from datetime import timedelta
 
 import pytest
 from django.utils import timezone
-from freezegun import freeze_time
 
 from sentry.dynamic_sampling.tasks.common import (
     GetActiveOrgs,
@@ -16,6 +15,7 @@ from sentry.dynamic_sampling.tasks.common import (
 from sentry.dynamic_sampling.tasks.task_context import DynamicSamplingLogState, TaskContext
 from sentry.snuba.metrics.naming_layer.mri import TransactionMRI
 from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase
+from sentry.testutils.helpers.datetime import freeze_time
 
 MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace(
     hour=0, minute=0, second=0, microsecond=0
@@ -47,7 +47,7 @@ class FakeContextIterator:
     def __next__(self):
         if self.count < 2:
             self.count += 1
-            self.frozen_time.tick(delta=timedelta(seconds=self.tick_seconds))
+            self.frozen_time.shift(self.tick_seconds)
             return self.count
         raise StopIteration()
 
@@ -186,10 +186,10 @@ def test_timed_function_correctly_times_inner_function():
         @timed_function()
         def f1(state: DynamicSamplingLogState, x: int, y: str):
             state.num_iterations = 1
-            frozen_time.tick()
+            frozen_time.shift(1)
 
         f1(context, 1, "x")
-        frozen_time.tick()
+        frozen_time.shift(1)
         f1(context, 1, "x")
 
         # two seconds passed inside f1 ( one for each call)
@@ -204,12 +204,12 @@ def test_timed_function_correctly_raises_when_task_expires():
         @timed_function()
         def f1(state: DynamicSamplingLogState, x: int, y: str):
             state.num_iterations = 1
-            frozen_time.tick()
+            frozen_time.shift(1)
 
         f1(context, 1, "x")
         t = context.get_timer("f1")
         assert t.current() == 1.0
-        frozen_time.tick()
+        frozen_time.shift(1)
         assert t.current() == 1.0  # timer should not be moving ouside the function
         f1(context, 1, "x")
 

+ 20 - 24
tests/sentry/dynamic_sampling/tasks/test_task_context.py

@@ -1,11 +1,7 @@
 import time
-from datetime import timedelta
-
-from freezegun import freeze_time
 
 from sentry.dynamic_sampling.tasks.task_context import DynamicSamplingLogState, TaskContext, Timers
-
-SECOND = timedelta(seconds=1)
+from sentry.testutils.helpers.datetime import freeze_time
 
 
 def test_task_context_expiration_time():
@@ -80,7 +76,7 @@ def test_timer_raw():
         t.stop()
 
         # jump 1 second in the future
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         # the timer is stopped so nothing should have happened
         assert t.current() == 0
@@ -96,7 +92,7 @@ def test_timer_raw():
         t.start()
 
         # another sec
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         # now the timer should be at 1 sec
         assert t.current() == 1.0
@@ -111,14 +107,14 @@ def test_timer_raw():
         t.stop()
         assert t.current() == 1.0
 
-        frozen_time.tick()
+        frozen_time.shift(1)
         # check that we can accumulate multiple stops and starts
         assert t.current() == 1.0
         t.start()
         assert t.current() == 1.0
 
         # another sec
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         assert t.current() == 2.0
         t.stop()
@@ -161,7 +157,7 @@ def test_named_timer_raw():
         ta.stop()
 
         # jump 1 second in the future
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         # the timer is stopped for a&b so nothing should have happened
         assert ta.current() == 0
@@ -184,7 +180,7 @@ def test_named_timer_raw():
         ta.start()
 
         # another sec
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         # now the timer should be at 1 sec
         assert ta.current() == 1.0
@@ -202,7 +198,7 @@ def test_named_timer_raw():
         assert tc.current() == 2.0
         tb.start()
 
-        frozen_time.tick()
+        frozen_time.shift(1)
         # check that we can accumulate multiple stops and starts
         assert ta.current() == 1.0
         assert tb.current() == 1.0
@@ -213,7 +209,7 @@ def test_named_timer_raw():
         assert tc.current() == 3.0
 
         # another sec
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         assert ta.current() == 2.0
         assert tb.current() == 2.0
@@ -224,7 +220,7 @@ def test_named_timer_raw():
         assert tc.current() == 4.0
 
         # another sec
-        frozen_time.tick()
+        frozen_time.shift(1)
 
         assert ta.current() == 2.0
         assert tb.current() == 3.0
@@ -243,8 +239,8 @@ def test_timer_context_manager():
                 assert t.current("a") == i
 
                 # jump one sec
-                frozen_time.tick()
-            frozen_time.tick()
+                frozen_time.shift(1)
+            frozen_time.shift(1)
 
         # we advanced 3 seconds within the timer and 3 outside we should have only counted 3
         assert t.current("a") == 3
@@ -265,21 +261,21 @@ def test_named_timer_context_manager():
 
                 with t.get_timer("a") as ta:
                     assert ta.current() == i
-                    frozen_time.tick(delta=SECOND)
+                    frozen_time.shift(1)
                 with t.get_timer("b") as tb:
                     assert tb.current() == i * 2
-                    frozen_time.tick(delta=SECOND * 2)
+                    frozen_time.shift(2)
                 with t.get_timer("c") as tc:
                     assert tc.current() == i * 3
-                    frozen_time.tick(delta=SECOND * 3)
+                    frozen_time.shift(3)
 
                 # jump one sec
-                frozen_time.tick(delta=SECOND)
+                frozen_time.shift(1)
 
-            frozen_time.tick(delta=SECOND)
+            frozen_time.shift(1)
 
         # outside the context manager timers should not advance
-        frozen_time.tick(delta=SECOND * 100)
+        frozen_time.shift(100)
 
         assert t.current("a") == 3
         assert t.current("b") == 3 * 2
@@ -292,10 +288,10 @@ def test_task_context_serialisation():
     with freeze_time("2023-07-12 10:00:00") as frozen_time:
         # a timer without state
         with task.get_timer("a"):
-            frozen_time.tick(delta=SECOND)
+            frozen_time.shift(1)
         # a timer with state
         with task.get_timer("b"):
-            frozen_time.tick(delta=SECOND * 2)
+            frozen_time.shift(2)
             state = task.get_function_state("b")
             state.num_iterations = 1
             state.num_orgs = 2

+ 3 - 6
tests/sentry/middleware/test_ratelimit_middleware.py

@@ -7,8 +7,6 @@ from django.contrib.auth.models import AnonymousUser
 from django.http.request import HttpRequest
 from django.test import RequestFactory, override_settings
 from django.urls import re_path, reverse
-from freezegun import freeze_time
-from freezegun.api import FrozenDateTimeFactory
 from rest_framework.permissions import AllowAny
 from rest_framework.response import Response
 
@@ -19,6 +17,7 @@ from sentry.models import ApiKey, ApiToken, SentryAppInstallation, User
 from sentry.ratelimits.config import RateLimitConfig, get_default_rate_limits_for_group
 from sentry.ratelimits.utils import get_rate_limit_config, get_rate_limit_key, get_rate_limit_value
 from sentry.testutils.cases import APITestCase, TestCase
+from sentry.testutils.helpers.datetime import freeze_time
 from sentry.testutils.silo import control_silo_test
 from sentry.types.ratelimit import RateLimit, RateLimitCategory
 
@@ -110,10 +109,9 @@ class RatelimitMiddlewareTest(TestCase):
         # Requests outside the current window should not be rate limited
         default_rate_limit_mock.return_value = RateLimit(1, 1)
         with freeze_time("2000-01-01") as frozen_time:
-            assert isinstance(frozen_time, FrozenDateTimeFactory)
             self.middleware.process_view(request, self._test_endpoint, [], {})
             assert not request.will_be_rate_limited
-            frozen_time.tick(1)
+            frozen_time.shift(1)
             self.middleware.process_view(request, self._test_endpoint, [], {})
             assert not request.will_be_rate_limited
 
@@ -128,10 +126,9 @@ class RatelimitMiddlewareTest(TestCase):
 
         default_rate_limit_mock.return_value = RateLimit(1, 1)
         with freeze_time("2000-01-01") as frozen_time:
-            assert isinstance(frozen_time, FrozenDateTimeFactory)
             self.middleware.process_view(request, self._test_endpoint, [], {})
             assert not request.will_be_rate_limited
-            frozen_time.tick(1)
+            frozen_time.shift(1)
             self.middleware.process_view(request, self._test_endpoint, [], {})
             assert not request.will_be_rate_limited
 

+ 4 - 4
tests/sentry/processing/realtime_metrics/test_redis.py

@@ -1,13 +1,13 @@
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import Any, Dict
 
 import pytest
-from freezegun import freeze_time
 from redis import StrictRedis
 
 from sentry.exceptions import InvalidConfiguration
 from sentry.processing import realtime_metrics
 from sentry.processing.realtime_metrics.redis import RedisRealtimeMetricsStore
+from sentry.testutils.helpers.datetime import freeze_time
 from sentry.utils import redis
 
 
@@ -60,7 +60,7 @@ def test_record_project_duration_same_bucket(
 ) -> None:
     with freeze_time(datetime.fromtimestamp(1147)) as frozen_datetime:
         store.record_project_duration(17, 1.0)
-        frozen_datetime.tick(delta=timedelta(seconds=2))
+        frozen_datetime.shift(2)
         store.record_project_duration(17, 1.0)
 
     assert redis_cluster.get("symbolicate_event_low_priority:budget:10:17:1140") == "2000"
@@ -71,7 +71,7 @@ def test_record_project_duration_different_buckets(
 ) -> None:
     with freeze_time(datetime.fromtimestamp(1147)) as frozen_datetime:
         store.record_project_duration(17, 1.0)
-        frozen_datetime.tick(delta=timedelta(seconds=5))
+        frozen_datetime.shift(5)
         store.record_project_duration(17, 1.0)
 
     assert redis_cluster.get("symbolicate_event_low_priority:budget:10:17:1140") == "1000"

+ 3 - 7
tests/sentry/ratelimits/test_redis.py

@@ -1,10 +1,8 @@
 from time import time
 
-from freezegun import freeze_time
-from freezegun.api import FrozenDateTimeFactory
-
 from sentry.ratelimits.redis import RedisRateLimiter
 from sentry.testutils.cases import TestCase
+from sentry.testutils.helpers.datetime import freeze_time
 from sentry.testutils.silo import region_silo_test
 
 
@@ -42,17 +40,15 @@ class RedisRateLimiterTest(TestCase):
     def test_current_value_expire(self):
         """Ensure that the count resets when the window expires"""
         with freeze_time("2000-01-01") as frozen_time:
-            assert isinstance(frozen_time, FrozenDateTimeFactory)
             for _ in range(10):
                 self.backend.is_limited("foo", 1, window=10)
             assert self.backend.current_value("foo", window=10) == 10
 
-            frozen_time.tick(10)
+            frozen_time.shift(10)
             assert self.backend.current_value("foo", window=10) == 0
 
     def test_is_limited_with_value(self):
         with freeze_time("2000-01-01") as frozen_time:
-            assert isinstance(frozen_time, FrozenDateTimeFactory)
             expected_reset_time = int(time() + 5)
 
             limited, value, reset_time = self.backend.is_limited_with_value("foo", 1, window=5)
@@ -65,7 +61,7 @@ class RedisRateLimiterTest(TestCase):
             assert value == 2
             assert reset_time == expected_reset_time
 
-            frozen_time.tick(5)
+            frozen_time.shift(5)
             limited, value, reset_time = self.backend.is_limited_with_value("foo", 1, window=5)
             assert not limited
             assert value == 1

+ 2 - 4
tests/sentry/rules/history/endpoints/test_project_rule_preview.py

@@ -2,11 +2,10 @@ from datetime import timedelta
 
 from dateutil.parser import parse as parse_datetime
 from django.utils import timezone
-from freezegun import freeze_time
-from freezegun.api import FrozenDateTimeFactory
 
 from sentry.models import Activity, Group, GroupInbox, GroupInboxReason
 from sentry.testutils.cases import APITestCase
+from sentry.testutils.helpers.datetime import freeze_time
 from sentry.testutils.silo import region_silo_test
 from sentry.types.activity import ActivityType
 
@@ -73,7 +72,6 @@ class ProjectRulePreviewEndpointTest(APITestCase):
     def test_endpoint(self):
         time_to_freeze = timezone.now()
         with freeze_time(time_to_freeze) as frozen_time:
-            assert isinstance(frozen_time, FrozenDateTimeFactory)
             resp = self.get_success_response(
                 self.organization.slug,
                 self.project.slug,
@@ -90,7 +88,7 @@ class ProjectRulePreviewEndpointTest(APITestCase):
             result = parse_datetime(resp["endpoint"])
             endpoint = time_to_freeze.replace(tzinfo=result.tzinfo)
             assert result == endpoint
-            frozen_time.tick(1)
+            frozen_time.shift(1)
 
             resp = self.get_success_response(
                 self.organization.slug,