def test_context_removed_after_exit(self): example_context = {"Hello": "World"} with set_current_context(example_context): pass with pytest.raises(AirflowException, ): get_current_context()
def test_nested_context(self): """ Nested execution context should be supported in case the user uses multiple context managers. Each time the execute method of an operator is called, we set a new 'current' context. This test verifies that no matter how many contexts are entered - order is preserved """ max_stack_depth = 15 ctx_list = [] for i in range(max_stack_depth): # Create all contexts in ascending order new_context = {"ContextId": i} # Like 15 nested with statements ctx_obj = set_current_context(new_context) ctx_obj.__enter__() # pylint: disable=E1101 ctx_list.append(ctx_obj) for i in reversed(range(max_stack_depth)): # Iterate over contexts in reverse order - stack is LIFO ctx = get_current_context() assert ctx["ContextId"] == i # End of with statement ctx_list[i].__exit__(None, None, None)
def test_current_context_roundtrip(self): example_context = {"Hello": "World"} with set_current_context(example_context): assert get_current_context() == example_context