From d04ab9367fc6789f1c26494adfde0c8cda92339a Mon Sep 17 00:00:00 2001
From: Trey Spiller <treyspiller@gmail.com>
Date: Fri, 13 Dec 2024 17:29:33 -0600
Subject: [PATCH] Update sqlmesh test freezegun call

---
 setup.py                        | 2 +-
 sqlmesh/core/test/definition.py | 6 ++++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index cb6db7f56..d7f63aa12 100644
--- 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
index ae62e9e85..bd41bfd58 100644
--- 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()