Switch tests from freezetime to time_machine library

TobikoData · Dec 13, 2024 · 96bd5a3 · 96bd5a3
1 parent 434be40
commit 96bd5a3
Show file tree

Hide file tree

Showing 12 changed files with 112 additions and 111 deletions.
diff --git a/.circleci/continue_config.yml b/.circleci/continue_config.yml
@@ -310,13 +310,13 @@ workflows:
                 - "3.11"
                 - "3.12"
       - style_and_slow_tests_pydantic_v1
-      - airflow_docker_tests:
-          requires:
-            - style_and_slow_tests
-          filters:
-            branches:
-              only:
-                - main
+      - airflow_docker_tests
+          # requires:
+          #   - style_and_slow_tests
+          # filters:
+          #   branches:
+          #     only:
+          #       - main
       - engine_tests_docker:
           name: engine_<< matrix.engine >>
           matrix:
@@ -334,8 +334,8 @@ workflows:
           name: cloud_engine_<< matrix.engine >>
           context:
             - sqlmesh_cloud_database_integration
-          requires:
-            - engine_tests_docker
+          # requires:
+          #   - engine_tests_docker
           matrix:
             parameters:
               engine:
@@ -345,10 +345,10 @@ workflows:
                 - bigquery
                 - clickhouse-cloud
                 - athena
-          filters:
-           branches:
-             only:
-               - main
+          # filters:
+          #  branches:
+          #    only:
+          #      - main
       - trigger_private_tests:
           requires:
             - style_and_slow_tests

diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py
@@ -6,7 +6,7 @@
 
 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
@@ -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

diff --git a/tests/core/test_context.py b/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"):
         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"):
         context = Context(paths=path, gateway="duckdb_persistent")
         expression = d.parse(
             """