chore: set airflow as library for some tests (#2265)

* chore: set airflow as library for some tests

* remove duplicate fixture

Author: eakmanrq

tests/schedulers/airflow/operators/fixtures.py

@@ -0,0 +1,10 @@
+import os
+from unittest import mock
+
+import pytest
+
+
+@pytest.fixture
+def set_airflow_as_library():
+    with mock.patch.dict(os.environ, {"_AIRFLOW__AS_LIBRARY": "1"}):
+        yield

tests/schedulers/airflow/operators/test_hwm_sensor.py

@@ -13,10 +13,11 @@
 )
 from sqlmesh.utils.date import to_datetime
 
+pytest_plugins = ["tests.schedulers.airflow.operators.fixtures"]
 pytestmark = pytest.mark.airflow
 
 
-def test_no_current_hwm(mocker: MockerFixture, make_snapshot, random_name):
+def test_no_current_hwm(mocker: MockerFixture, make_snapshot, random_name, set_airflow_as_library):
     this_snapshot = make_snapshot(SqlModel(name="this", query=parse_one("select 1, ds")))
     this_snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
@@ -43,7 +44,7 @@ def test_no_current_hwm(mocker: MockerFixture, make_snapshot, random_name):
     get_snapshots_mock.assert_called_once_with([target_snapshot.table_info])
 
 
-def test_current_hwm_below_target(mocker: MockerFixture, make_snapshot):
+def test_current_hwm_below_target(mocker: MockerFixture, make_snapshot, set_airflow_as_library):
     this_snapshot = make_snapshot(
         SqlModel(name="this", query=parse_one("select 1, ds")), version="a"
     )
@@ -82,7 +83,7 @@ def test_current_hwm_below_target(mocker: MockerFixture, make_snapshot):
     get_snapshots_mock.assert_called_once_with([target_snapshot_v1.table_info])
 
 
-def test_current_hwm_above_target(mocker: MockerFixture, make_snapshot):
+def test_current_hwm_above_target(mocker: MockerFixture, make_snapshot, set_airflow_as_library):
     this_snapshot = make_snapshot(
         SqlModel(name="this", query=parse_one("select 1, ds")), version="a"
     )
@@ -115,7 +116,7 @@ def test_current_hwm_above_target(mocker: MockerFixture, make_snapshot):
     get_snapshots_mock.assert_called_once_with([target_snapshot_v1.table_info])
 
 
-def test_hwm_external_sensor(mocker: MockerFixture, make_snapshot):
+def test_hwm_external_sensor(mocker: MockerFixture, make_snapshot, set_airflow_as_library):
     snapshot = make_snapshot(
         SqlModel(
             name="this",

tests/schedulers/airflow/operators/test_targets.py

@@ -18,6 +18,7 @@
 from sqlmesh.schedulers.airflow.operators import targets
 from sqlmesh.utils.date import to_datetime
 
+pytest_plugins = ["tests.schedulers.airflow.operators.fixtures"]
 pytestmark = pytest.mark.airflow
 
 
@@ -29,7 +30,9 @@ def model() -> Model:
     )
 
 
-def test_evaluation_target_execute(mocker: MockerFixture, make_snapshot: t.Callable, model: Model):
+def test_evaluation_target_execute(
+    mocker: MockerFixture, make_snapshot: t.Callable, model: Model, set_airflow_as_library
+):
     interval_ds = to_datetime("2022-01-01")
     logical_ds = to_datetime("2022-01-02")
 
@@ -76,7 +79,9 @@ def test_evaluation_target_execute(mocker: MockerFixture, make_snapshot: t.Calla
     )
 
 
-def test_evaluation_target_execute_seed_model(mocker: MockerFixture, make_snapshot: t.Callable):
+def test_evaluation_target_execute_seed_model(
+    mocker: MockerFixture, make_snapshot: t.Callable, set_airflow_as_library
+):
     interval_ds = to_datetime("2022-01-01")
     logical_ds = to_datetime("2022-01-02")
 
@@ -135,7 +140,9 @@ def test_evaluation_target_execute_seed_model(mocker: MockerFixture, make_snapsh
     )
 
 
-def test_cleanup_target_execute(mocker: MockerFixture, make_snapshot: t.Callable, model: Model):
+def test_cleanup_target_execute(
+    mocker: MockerFixture, make_snapshot: t.Callable, model: Model, set_airflow_as_library
+):
     snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
 
@@ -181,7 +188,7 @@ def test_cleanup_target_execute(mocker: MockerFixture, make_snapshot: t.Callable
 
 
 def test_cleanup_target_skip_execution(
-    mocker: MockerFixture, make_snapshot: t.Callable, model: Model
+    mocker: MockerFixture, make_snapshot: t.Callable, model: Model, set_airflow_as_library
 ):
     snapshot = make_snapshot(model)
     snapshot.version = "test_version"