44
44
from airflow .config_templates .airflow_local_settings import DEFAULT_LOGGING_CONFIG
45
45
from airflow .exceptions import (
46
46
AirflowException ,
47
+ AirflowProviderDeprecationWarning ,
47
48
DeserializingResultError ,
48
49
)
49
50
from airflow .models .baseoperator import BaseOperator
@@ -1810,22 +1811,35 @@ def default_kwargs(*, python_version=DEFAULT_PYTHON_VERSION, **kwargs):
1810
1811
1811
1812
class TestCurrentContext :
1812
1813
def test_current_context_no_context_raise (self ):
1813
- with pytest .raises (RuntimeError ):
1814
- get_current_context ()
1814
+ if AIRFLOW_V_3_0_PLUS :
1815
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1816
+ with pytest .raises (RuntimeError ):
1817
+ get_current_context ()
1818
+ else :
1819
+ with pytest .raises (RuntimeError ):
1820
+ get_current_context ()
1815
1821
1816
1822
def test_current_context_roundtrip (self ):
1817
1823
example_context = {"Hello" : "World" }
1818
-
1819
1824
with set_current_context (example_context ):
1820
- assert get_current_context () == example_context
1825
+ if AIRFLOW_V_3_0_PLUS :
1826
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1827
+ assert get_current_context () == example_context
1828
+ else :
1829
+ assert get_current_context () == example_context
1821
1830
1822
1831
def test_context_removed_after_exit (self ):
1823
1832
example_context = {"Hello" : "World" }
1824
1833
1825
1834
with set_current_context (example_context ):
1826
1835
pass
1827
- with pytest .raises (RuntimeError ):
1828
- get_current_context ()
1836
+ if AIRFLOW_V_3_0_PLUS :
1837
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1838
+ with pytest .raises (RuntimeError ):
1839
+ get_current_context ()
1840
+ else :
1841
+ with pytest .raises (RuntimeError ):
1842
+ get_current_context ()
1829
1843
1830
1844
def test_nested_context (self ):
1831
1845
"""
@@ -1842,12 +1856,21 @@ def test_nested_context(self):
1842
1856
ctx_obj = set_current_context (new_context )
1843
1857
ctx_obj .__enter__ ()
1844
1858
ctx_list .append (ctx_obj )
1845
- for i in reversed (range (max_stack_depth )):
1846
- # Iterate over contexts in reverse order - stack is LIFO
1847
- ctx = get_current_context ()
1848
- assert ctx ["ContextId" ] == i
1849
- # End of with statement
1850
- ctx_list [i ].__exit__ (None , None , None )
1859
+ if AIRFLOW_V_3_0_PLUS :
1860
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1861
+ for i in reversed (range (max_stack_depth )):
1862
+ # Iterate over contexts in reverse order - stack is LIFO
1863
+ ctx = get_current_context ()
1864
+ assert ctx ["ContextId" ] == i
1865
+ # End of with statement
1866
+ ctx_list [i ].__exit__ (None , None , None )
1867
+ else :
1868
+ for i in reversed (range (max_stack_depth )):
1869
+ # Iterate over contexts in reverse order - stack is LIFO
1870
+ ctx = get_current_context ()
1871
+ assert ctx ["ContextId" ] == i
1872
+ # End of with statement
1873
+ ctx_list [i ].__exit__ (None , None , None )
1851
1874
1852
1875
1853
1876
class MyContextAssertOperator (BaseOperator ):
@@ -1889,12 +1912,20 @@ class TestCurrentContextRuntime:
1889
1912
def test_context_in_task (self ):
1890
1913
with DAG (dag_id = "assert_context_dag" , default_args = DEFAULT_ARGS , schedule = "@once" ):
1891
1914
op = MyContextAssertOperator (task_id = "assert_context" )
1892
- op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1915
+ if AIRFLOW_V_3_0_PLUS :
1916
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1917
+ op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1918
+ else :
1919
+ op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1893
1920
1894
1921
def test_get_context_in_old_style_context_task (self ):
1895
1922
with DAG (dag_id = "edge_case_context_dag" , default_args = DEFAULT_ARGS , schedule = "@once" ):
1896
1923
op = PythonOperator (python_callable = get_all_the_context , task_id = "get_all_the_context" )
1897
- op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1924
+ if AIRFLOW_V_3_0_PLUS :
1925
+ with pytest .warns (AirflowProviderDeprecationWarning ):
1926
+ op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1927
+ else :
1928
+ op .run (ignore_first_depends_on_past = True , ignore_ti_state = True )
1898
1929
1899
1930
1900
1931
@pytest .mark .need_serialized_dag (False )
0 commit comments