add test

Signed-off-by: Katrina Rogan <[email protected]>
flyteorg · Mar 29, 2021 · 258a54c · 258a54c
1 parent a37bab1
commit 258a54c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/flytekit/core/context_manager.py b/flytekit/core/context_manager.py
@@ -415,7 +415,7 @@ def new_execution_context(
         if additional_context is None:
             additional_context = self.execution_state.additional_context if self.execution_state is not None else None
         elif self.execution_state is not None and self.execution_state.additional_context is not None:
-            additional_context = {**additional_context, **self.execution_state.additional_context}
+            additional_context = {**self.execution_state.additional_context, **additional_context}
         exec_state = ExecutionState(
             mode=mode, working_dir=working_dir, engine_dir=engine_dir, additional_context=additional_context
         )

diff --git a/tests/flytekit/unit/core/test_context_manager.py b/tests/flytekit/unit/core/test_context_manager.py
@@ -1,4 +1,4 @@
-from flytekit.core.context_manager import CompilationState, FlyteContext, look_up_image_info
+from flytekit.core.context_manager import CompilationState, ExecutionState, FlyteContext, look_up_image_info
 
 
 class SampleTestClass(object):
@@ -42,3 +42,14 @@ def test_look_up_image_info():
     assert img.name == "x"
     assert img.tag == "latest"
     assert img.fqn == "localhost:5000/xyz"
+
+
+def test_additional_context():
+    with FlyteContext.current_context() as ctx:
+        with ctx.new_execution_context(
+            mode=ExecutionState.Mode.TASK_EXECUTION, additional_context={1: "outer", 2: "foo"}
+        ) as exec_ctx_outer:
+            with exec_ctx_outer.new_execution_context(
+                mode=ExecutionState.Mode.TASK_EXECUTION, additional_context={1: "inner", 3: "baz"}
+            ) as exec_ctx_inner:
+                assert exec_ctx_inner.execution_state.additional_context == {1: "inner", 2: "foo", 3: "baz"}