Update sqlmesh test freezegun call

TobikoData · Dec 13, 2024 · d04ab93 · d04ab93
1 parent d2bf20d
commit d04ab93
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -37,7 +37,6 @@
         "croniter",
         "duckdb!=0.10.3",
         "dateparser",
-        "freezegun",
         "hyperscript>=0.1.0",
         "importlib-metadata; python_version<'3.12'",
         "ipywidgets",
@@ -50,6 +49,7 @@
         "setuptools; python_version>='3.12'",
         "sqlglot[rs]~=26.0.0",
         "tenacity",
+        "time-machine",
     ],
     extras_require={
         "athena": ["PyAthena[Pandas]"],

diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
@@ -12,7 +12,7 @@
 import numpy as np
 import pandas as pd
 from io import StringIO
-from freezegun import freeze_time
+import time_machine
 from pandas.api.types import is_object_dtype
 from sqlglot import Dialect, exp
 from sqlglot.optimizer.annotate_types import annotate_types
@@ -649,7 +649,9 @@ def runTest(self) -> None:
 
     def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
-        time_ctx = freeze_time(self._execution_time) if self._execution_time else nullcontext()
+        time_ctx = (
+            time_machine.travel(self._execution_time) if self._execution_time else nullcontext()
+        )
         with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
             with t.cast(AbstractContextManager, time_ctx):
                 variables = self.body.get("vars", {}).copy()