Chore: switch tests from freezegun to time-machine

setup.py

@@ -94,6 +94,7 @@
             "ruff~=0.7.0",
             "snowflake-connector-python[pandas,secure-local-storage]>=3.0.2",
             "sqlalchemy-stubs",
+            "time-machine",
             "types-croniter",
             "types-dateparser",
             "types-python-dateutil",

tests/cli/test_cli.py

@@ -6,15 +6,15 @@
 
 import pytest
 from click.testing import CliRunner
-from freezegun import freeze_time
+import time_machine
 
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 from sqlmesh.cli.main import cli
 from sqlmesh.core.context import Context
 from sqlmesh.integrations.dlt import generate_dlt_models
 from sqlmesh.utils.date import yesterday_ds
 
-FREEZE_TIME = "2023-01-01 00:00:00"
+FREEZE_TIME = "2023-01-01 00:00:00 UTC"
 
 pytestmark = pytest.mark.slow
 
@@ -658,7 +658,7 @@ def test_run_no_prod(runner, tmp_path):
 
 
 @pytest.mark.parametrize("flag", ["--skip-backfill", "--dry-run"])
-@freeze_time(FREEZE_TIME)
+@time_machine.travel(FREEZE_TIME)
 def test_run_dev(runner, tmp_path, flag):
     create_example_project(tmp_path)
 
@@ -676,7 +676,7 @@ def test_run_dev(runner, tmp_path, flag):
     assert_model_batches_executed(result)
 
 
-@freeze_time(FREEZE_TIME)
+@time_machine.travel(FREEZE_TIME)
 def test_run_cron_not_elapsed(runner, tmp_path, caplog):
     create_example_project(tmp_path)
     init_prod_and_backfill(runner, tmp_path)
@@ -692,11 +692,12 @@ def test_run_cron_elapsed(runner, tmp_path):
     create_example_project(tmp_path)
 
     # Create and backfill `prod` environment
-    with freeze_time("2023-01-01 23:59:00"):
+    with time_machine.travel("2023-01-01 23:59:00 UTC", tick=False) as traveler:
+        runner = CliRunner()
         init_prod_and_backfill(runner, tmp_path)
 
-    # Run `prod` environment with daily cron elapsed
-    with freeze_time("2023-01-02 00:01:00"):
+        # Run `prod` environment with daily cron elapsed
+        traveler.move_to("2023-01-02 00:01:00 UTC")
         result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "run"])
 
     assert result.exit_code == 0

tests/core/test_context.py

@@ -6,7 +6,7 @@
 from tempfile import TemporaryDirectory
 from unittest.mock import PropertyMock, call, patch
 
-import freezegun
+import time_machine
 import pytest
 import pandas as pd
 from pathlib import Path
@@ -850,15 +850,15 @@ def test_plan_default_end(sushi_context_pre_scheduling: Context):
 @pytest.mark.slow
 def test_plan_start_ahead_of_end(copy_to_temp_path):
     path = copy_to_temp_path("examples/sushi")
-    with freezegun.freeze_time("2024-01-02 00:00:00"):
+    with time_machine.travel("2024-01-02 00:00:00 UTC"):
         context = Context(paths=path, gateway="duckdb_persistent")
         context.plan("prod", no_prompts=True, auto_apply=True)
         assert all(
             i == to_timestamp("2024-01-02")
             for i in context.state_sync.max_interval_end_per_model("prod").values()
         )
         context.close()
-    with freezegun.freeze_time("2024-01-03 00:00:00"):
+    with time_machine.travel("2024-01-03 00:00:00 UTC"):
         context = Context(paths=path, gateway="duckdb_persistent")
         expression = d.parse(
             """