This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 584050de80f Replace `mock.patch("utcnow")` with time_machine. (#53642)
584050de80f is described below
commit 584050de80f4d818acd3298539c4c532f9b7d75d
Author: Ash Berlin-Taylor <[email protected]>
AuthorDate: Tue Jul 22 21:16:16 2025 +0100
Replace `mock.patch("utcnow")` with time_machine. (#53642)
As part of a future PR we are going to move the location of this function,
and more
generally we shouldn't be mocking a specific function, time_machine is
better
suited to this task.
It would be nice if we could add a ruff/etc rule to forbid this, but I can't
think of a way of enforcing this.
---
.../utils/test_exponential_backoff_retry.py | 56 ++++++++++------------
.../unit/databricks/utils/test_openlineage.py | 30 +++++-------
.../unit/openlineage/plugins/test_listener.py | 25 +++++-----
.../tests/unit/snowflake/utils/test_openlineage.py | 25 ++++------
.../tests/unit/standard/utils/test_skipmixin.py | 7 ++-
5 files changed, 63 insertions(+), 80 deletions(-)
diff --git
a/providers/amazon/tests/unit/amazon/aws/executors/utils/test_exponential_backoff_retry.py
b/providers/amazon/tests/unit/amazon/aws/executors/utils/test_exponential_backoff_retry.py
index c180184c30d..e8ff928d4ac 100644
---
a/providers/amazon/tests/unit/amazon/aws/executors/utils/test_exponential_backoff_retry.py
+++
b/providers/amazon/tests/unit/amazon/aws/executors/utils/test_exponential_backoff_retry.py
@@ -16,7 +16,7 @@
# under the License.
from __future__ import annotations
-from datetime import datetime
+from datetime import datetime, timezone
from unittest import mock
import pytest
@@ -28,18 +28,16 @@ from
airflow.providers.amazon.aws.executors.utils.exponential_backoff_retry impo
class TestExponentialBackoffRetry:
- @mock.patch("airflow.utils.timezone.utcnow")
- def test_exponential_backoff_retry_base_case(self, mock_utcnow):
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 0, 5)
+ def test_exponential_backoff_retry_base_case(self, time_machine):
+ time_machine.move_to(datetime(2023, 1, 1, 12, 0, 5))
mock_callable_function = mock.Mock()
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=0,
callable_function=mock_callable_function,
)
mock_callable_function.assert_called_once()
- @mock.patch("airflow.utils.timezone.utcnow")
@pytest.mark.parametrize(
"attempt_number, utcnow_value, expected_calls",
[
@@ -111,28 +109,27 @@ class TestExponentialBackoffRetry:
],
)
def test_exponential_backoff_retry_parameterized(
- self, mock_utcnow, attempt_number, utcnow_value, expected_calls
+ self, attempt_number, utcnow_value, expected_calls, time_machine
):
+ time_machine.move_to(utcnow_value)
mock_callable_function = mock.Mock()
mock_callable_function.__name__ = "test_callable_function"
mock_callable_function.side_effect = Exception()
- mock_utcnow.return_value = utcnow_value
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=attempt_number,
callable_function=mock_callable_function,
)
assert mock_callable_function.call_count == expected_calls
- @mock.patch("airflow.utils.timezone.utcnow")
- def test_exponential_backoff_retry_fail_success(self, mock_utcnow, caplog):
+ def test_exponential_backoff_retry_fail_success(self, time_machine,
caplog):
mock_callable_function = mock.Mock()
mock_callable_function.__name__ = "test_callable_function"
mock_callable_function.side_effect = [Exception(), True]
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 0, 2)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 0, 2))
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=0,
callable_function=mock_callable_function,
)
@@ -140,39 +137,38 @@ class TestExponentialBackoffRetry:
assert any("Error calling" in log for log in caplog.messages)
caplog.clear() # clear messages so that we have clean logs for the
next call
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 0, 6)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 0, 6))
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=1,
callable_function=mock_callable_function,
)
assert all("Error calling" not in log for log in caplog.messages)
- @mock.patch("airflow.utils.timezone.utcnow")
- def test_exponential_backoff_retry_max_delay(self, mock_utcnow):
+ def test_exponential_backoff_retry_max_delay(self, time_machine):
mock_callable_function = mock.Mock()
mock_callable_function.__name__ = "test_callable_function"
mock_callable_function.return_value = Exception()
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 4, 15)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 4, 15))
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=4,
callable_function=mock_callable_function,
max_delay=60 * 5,
)
mock_callable_function.assert_not_called() # delay is 256 seconds; no
calls made
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 4, 16)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 4, 16))
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=4,
callable_function=mock_callable_function,
max_delay=60 * 5,
)
mock_callable_function.assert_called_once()
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 5, 0)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 5, 0))
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=5,
callable_function=mock_callable_function,
max_delay=60 * 5,
@@ -180,15 +176,14 @@ class TestExponentialBackoffRetry:
# delay should be 4^5=1024 seconds, but max_delay is 60*5=300 seconds
assert mock_callable_function.call_count == 2
- @mock.patch("airflow.utils.timezone.utcnow")
- def test_exponential_backoff_retry_max_attempts(self, mock_utcnow, caplog):
+ def test_exponential_backoff_retry_max_attempts(self, time_machine,
caplog):
mock_callable_function = mock.Mock()
mock_callable_function.__name__ = "test_callable_function"
mock_callable_function.return_value = Exception()
- mock_utcnow.return_value = datetime(2023, 1, 1, 12, 55, 0)
+ time_machine.move_to(datetime(2023, 1, 1, 12, 55, 0))
for i in range(10):
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=i,
callable_function=mock_callable_function,
max_attempts=3,
@@ -196,7 +191,6 @@ class TestExponentialBackoffRetry:
assert any("Max attempts reached." in log for log in caplog.messages)
assert mock_callable_function.call_count == 3
- @mock.patch("airflow.utils.timezone.utcnow")
@pytest.mark.parametrize(
"attempt_number, utcnow_value, expected_calls",
[
@@ -268,15 +262,15 @@ class TestExponentialBackoffRetry:
],
)
def test_exponential_backoff_retry_exponent_base_parameterized(
- self, mock_utcnow, attempt_number, utcnow_value, expected_calls
+ self, time_machine, attempt_number, utcnow_value, expected_calls
):
mock_callable_function = mock.Mock()
mock_callable_function.__name__ = "test_callable_function"
mock_callable_function.side_effect = Exception()
- mock_utcnow.return_value = utcnow_value
+ time_machine.move_to(utcnow_value)
exponential_backoff_retry(
- last_attempt_time=datetime(2023, 1, 1, 12, 0, 0),
+ last_attempt_time=datetime(2023, 1, 1, 12, 0, 0,
tzinfo=timezone.utc),
attempts_since_last_successful=attempt_number,
callable_function=mock_callable_function,
exponent_base=3,
diff --git
a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
index 24a8f96df25..20305a9af4e 100644
--- a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
+++ b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
@@ -285,13 +285,12 @@ def test_create_ol_event_pair_success(mock_generate_uuid,
is_successful):
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
-def test_emit_openlineage_events_for_databricks_queries(mock_now,
mock_generate_uuid, mock_version):
+def test_emit_openlineage_events_for_databricks_queries(mock_generate_uuid,
mock_version, time_machine):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1", "query2", "query3"]
original_query_ids = copy.deepcopy(query_ids)
@@ -523,15 +522,14 @@ def
test_emit_openlineage_events_for_databricks_queries(mock_now, mock_generate_
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def test_emit_openlineage_events_for_databricks_queries_without_metadata(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1"]
original_query_ids = copy.deepcopy(query_ids)
@@ -642,15 +640,14 @@ def
test_emit_openlineage_events_for_databricks_queries_without_metadata(
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1"]
hook = mock.MagicMock()
@@ -765,15 +762,14 @@ def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_i
)
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids_and_namespace(
- mock_now, mock_generate_uuid, mock_version, mock_parser
+ mock_generate_uuid, mock_version, mock_parser, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1"]
hook = mock.MagicMock()
@@ -884,15 +880,14 @@ def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_i
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids_and_namespace_raw_ns(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1"]
hook = DatabricksHook()
@@ -1004,15 +999,14 @@ def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_i
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_databricks_queries_ith_query_ids_and_hook_query_ids(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
hook = DatabricksSqlHook()
hook.query_ids = ["query2", "query3"]
diff --git
a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py
b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py
index 2f171a6b44b..bd57141a5a6 100644
--- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py
+++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py
@@ -420,16 +420,15 @@ class TestOpenLineageListenerAirflow2:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def test_adapter_fail_task_is_called_with_proper_arguments(
self,
- mock_utcnow,
mock_get_user_provided_run_facets,
mock_get_airflow_run_facet,
mock_get_task_parent_run_facet,
mock_disabled,
mock_debug_facet,
mock_debug_mode,
+ time_machine,
):
"""Tests that the 'fail_task' method of the OpenLineageAdapter is
invoked with the correct arguments.
@@ -438,6 +437,7 @@ class TestOpenLineageListenerAirflow2:
the test verifies the integrity and consistency of the data passed to
the adapter during task
failure events, thus confirming that the adapter's failure handling is
functioning as expected.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1))
listener, task_instance = self._create_listener_and_task_instance()
task_instance.logical_date = timezone.datetime(2020, 1, 1, 1, 1, 1)
@@ -551,16 +551,15 @@ class TestOpenLineageListenerAirflow2:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def test_adapter_complete_task_is_called_with_proper_arguments(
self,
- mock_utcnow,
mock_get_user_provided_run_facets,
mock_get_airflow_run_facet,
mock_get_task_parent_run_facet,
mock_disabled,
mock_debug_facet,
mock_debug_mode,
+ time_machine,
):
"""Tests that the 'complete_task' method of the OpenLineageAdapter is
called with the correct arguments.
@@ -570,6 +569,7 @@ class TestOpenLineageListenerAirflow2:
accordingly. This helps confirm the consistency and correctness of the
data passed to the adapter
during the task's lifecycle events.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1))
listener, task_instance = self._create_listener_and_task_instance()
mock_get_user_provided_run_facets.return_value = {"custom_user_facet":
2}
@@ -1301,16 +1301,15 @@ class TestOpenLineageListenerAirflow3:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def test_adapter_fail_task_is_called_with_proper_arguments(
self,
- mock_utcnow,
mock_get_user_provided_run_facets,
mock_get_airflow_run_facet,
mock_get_task_parent_run_facet,
mock_disabled,
mock_debug_facet,
mock_debug_mode,
+ time_machine,
):
"""Tests that the 'fail_task' method of the OpenLineageAdapter is
invoked with the correct arguments.
@@ -1319,6 +1318,7 @@ class TestOpenLineageListenerAirflow3:
the test verifies the integrity and consistency of the data passed to
the adapter during task
failure events, thus confirming that the adapter's failure handling is
functioning as expected.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1),
tick=False)
listener, task_instance = self._create_listener_and_task_instance()
task_instance.get_template_context()["dag_run"].logical_date =
timezone.datetime(2020, 1, 1, 1, 1, 1)
@@ -1426,20 +1426,20 @@ class TestOpenLineageListenerAirflow3:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def
test_adapter_fail_task_is_called_with_proper_arguments_for_db_task_instance_model(
self,
- mock_utcnow,
mock_get_task_parent_run_facet,
mock_debug_facet,
mock_debug_mode,
mock_emit,
+ time_machine,
):
"""Tests that the 'fail_task' method of the OpenLineageAdapter is
invoked with the correct arguments.
This particular test is using TaskInstance model available on API
Server and not on worker,
to simulate the listener being called after task's state has been
manually set via API.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1),
tick=False)
listener, task_instance =
self._create_listener_and_task_instance(runtime_ti=False)
mock_get_task_parent_run_facet.return_value = {"parent": 4}
@@ -1484,16 +1484,15 @@ class TestOpenLineageListenerAirflow3:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def test_adapter_complete_task_is_called_with_proper_arguments(
self,
- mock_utcnow,
mock_get_user_provided_run_facets,
mock_get_airflow_run_facet,
mock_get_task_parent_run_facet,
mock_disabled,
mock_debug_facet,
mock_debug_mode,
+ time_machine,
):
"""Tests that the 'complete_task' method of the OpenLineageAdapter is
called with the correct arguments.
@@ -1503,6 +1502,7 @@ class TestOpenLineageListenerAirflow3:
accordingly. This helps confirm the consistency and correctness of the
data passed to the adapter
during the task's lifecycle events.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1),
tick=False)
listener, task_instance = self._create_listener_and_task_instance()
mock_get_user_provided_run_facets.return_value = {"custom_user_facet":
2}
@@ -1607,15 +1607,16 @@ class TestOpenLineageListenerAirflow3:
@mock.patch(
"airflow.providers.openlineage.plugins.listener.OpenLineageListener._execute",
new=regular_call
)
- @mock.patch("airflow.utils.timezone.utcnow",
return_value=timezone.datetime(2023, 1, 3, 13, 1, 1))
def
test_adapter_complete_task_is_called_with_proper_arguments_for_db_task_instance_model(
- self, mock_utcnow, mock_get_task_parent_run_facet, mock_debug_facet,
mock_debug_mode, mock_emit
+ self, mock_get_task_parent_run_facet, mock_debug_facet,
mock_debug_mode, mock_emit, time_machine
):
"""Tests that the 'complete_task' method of the OpenLineageAdapter is
called with the correct arguments.
This particular test is using TaskInstance model available on API
Server and not on worker,
to simulate the listener being called after task's state has been
manually set via API.
"""
+ time_machine.move_to(timezone.datetime(2023, 1, 3, 13, 1, 1),
tick=False)
+
listener, task_instance =
self._create_listener_and_task_instance(runtime_ti=False)
mock_get_task_parent_run_facet.return_value = {"parent": 4}
mock_debug_facet.return_value = {"debug": "packages"}
diff --git a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
index dd9820de21c..4873210d2f0 100644
--- a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
+++ b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
@@ -580,15 +580,14 @@ def
test_create_snowflake_event_pair_success(mock_generate_uuid, is_successful):
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def test_emit_openlineage_events_for_snowflake_queries_with_extra_metadata(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1", "query2", "query3"]
original_query_ids = copy.deepcopy(query_ids)
@@ -821,15 +820,14 @@ def
test_emit_openlineage_events_for_snowflake_queries_with_extra_metadata(
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def test_emit_openlineage_events_for_snowflake_queries_without_extra_metadata(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
query_ids = ["query1"]
original_query_ids = copy.deepcopy(query_ids)
@@ -940,15 +938,14 @@ def
test_emit_openlineage_events_for_snowflake_queries_without_extra_metadata(
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def test_emit_openlineage_events_for_snowflake_queries_without_query_ids(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
hook = mock.MagicMock()
hook.query_ids = ["query1"]
@@ -1061,15 +1058,14 @@ def
test_emit_openlineage_events_for_snowflake_queries_without_query_ids(
@mock.patch("airflow.providers.openlineage.sqlparser.SQLParser.create_namespace",
return_value="snowflake_ns")
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_snowflake_queries_without_query_ids_and_namespace(
- mock_now, mock_generate_uuid, mock_version, mock_parser
+ mock_generate_uuid, mock_version, mock_parser, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
hook = mock.MagicMock()
hook.query_ids = ["query1"]
@@ -1181,15 +1177,14 @@ def
test_emit_openlineage_events_for_snowflake_queries_without_query_ids_and_nam
@mock.patch("importlib.metadata.version", return_value="2.3.0")
@mock.patch("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.utils.timezone.utcnow")
def
test_emit_openlineage_events_for_snowflake_queries_with_query_ids_and_hook_query_ids(
- mock_now, mock_generate_uuid, mock_version
+ mock_generate_uuid, mock_version, time_machine
):
fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
mock_generate_uuid.return_value = fake_uuid
default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- mock_now.return_value = default_event_time
+ time_machine.move_to(default_event_time, tick=False)
hook = mock.MagicMock()
hook.query_ids = ["query1"]
diff --git a/providers/standard/tests/unit/standard/utils/test_skipmixin.py
b/providers/standard/tests/unit/standard/utils/test_skipmixin.py
index 7d4c9875efb..4c79f3f893c 100644
--- a/providers/standard/tests/unit/standard/utils/test_skipmixin.py
+++ b/providers/standard/tests/unit/standard/utils/test_skipmixin.py
@@ -18,7 +18,7 @@
from __future__ import annotations
import datetime
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import MagicMock, Mock
import pytest
@@ -60,10 +60,9 @@ class TestSkipMixin:
def teardown_method(self):
self.clean_db()
- @patch("airflow.utils.timezone.utcnow")
- def test_skip(self, mock_now, dag_maker, session):
+ def test_skip(self, dag_maker, session, time_machine):
now = datetime.datetime.now(tz=datetime.timezone.utc)
- mock_now.return_value = now
+ time_machine.move_to(now, tick=False)
with dag_maker("dag"):
tasks = [EmptyOperator(task_id="task")]
