Repository: PrefectHQ/prefect

Issue:
Prefect context manager overwrites global contexts only once
## Description
<!-- A clear description of the bug -->

The [document](https://docs.prefect.io/core/concepts/execution.html#modifying-context-at-runtime) states we can modify global context keys by using a context manager:

> ### Modifying context at runtime
> Modifying context, even globally set context keys, at specific times is possible using a provided context manager:

However, in the attached script, the context key `foo` is overwritten only once.

## Expected Behavior
<!-- What did you expect to happen instead? -->

Modified context values do not change during a flow run.

## Reproduction
<!-- A minimal example that exhibits the behavior. -->

### context-test.py
```python
import prefect
from prefect import task, Flow

@task
def print_contexts():
    print("foo = " + prefect.context.foo)

with Flow('test') as flow:
    t1 = print_contexts()
    t2 = print_contexts()
    t1.set_downstream(t2)

with prefect.context(foo='foo'):
    flow.run()
```

### ~/.prefect/config.toml
```
[context]
foo = "global"
```

### Execution result
```
% python context-test.py
[2020-12-14 16:12:46+0900] INFO - prefect.FlowRunner | Beginning Flow run for 'test'
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Starting task run...
foo = foo
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Finished task run for task with final state: 'Success'
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Starting task run...
foo = global  <-- UNEXPECTED!!
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Finished task run for task with final state: 'Success'
[2020-12-14 16:12:46+0900] INFO - prefect.FlowRunner | Flow run SUCCESS: all reference tasks succeeded
```

## Environment
<!-- Any additional information about your environment

Optionally run `prefect diagnostics` from the command line and paste the information here. -->

```
% prefect diagnostics
{
  "config_overrides": {
    "context": {
      "foo": true
    }
  },
  "env_vars": [],
  "system_information": {
    "platform": "macOS-10.16-x86_64-i386-64bit",
    "prefect_backend": "cloud",
    "prefect_version": "0.13.19",
    "python_version": "3.8.5"
  }
}
```

Comments:

Patch:
diff --git a/changes/pr3924.yaml b/changes/pr3924.yaml
new file mode 100644
index 0000000000..de3d6b199d
--- /dev/null
+++ b/changes/pr3924.yaml
@@ -0,0 +1,2 @@
+fix:
+  - "Fix bug in `prefect.context` contextmanager that resulted in context fields reverting to their initially configured values - [#3924](https://github.com/PrefectHQ/prefect/pull/3924)"
diff --git a/src/prefect/utilities/context.py b/src/prefect/utilities/context.py
index d2bde4720e..27f4544bc6 100644
--- a/src/prefect/utilities/context.py
+++ b/src/prefect/utilities/context.py
@@ -82,10 +82,14 @@ class Context(DotDict, threading.local):
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        if "context" in config:
-            self.update(config.context)
-        self["config"] = merge_dicts(config, self.get("config", {}))  # order matters
+        init = {}
+        # Initialize with config context
+        init.update(config.get("context", {}))
+        # Overwrite with explicit args
+        init.update(dict(*args, **kwargs))
+        # Merge in config (with explicit args overwriting)
+        init["config"] = merge_dicts(config, init.get("config", {}))
+        super().__init__(init)
 
     def __getstate__(self) -> None:
         """
@@ -111,7 +115,8 @@ class Context(DotDict, threading.local):
             with prefect.context(dict(a=1, b=2), c=3):
                 print(prefect.context.a) # 1
         """
-        previous_context = self.copy()
+        # Avoid creating new `Context` object, copy as `dict` instead.
+        previous_context = self.__dict__.copy()
         try:
             new_context = dict(*args, **kwargs)
             if "config" in new_context:


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
index 79f7a2c0f6..b17a2a3fc0 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -120,11 +120,13 @@ def test_modify_context_by_calling_update_inside_contextmanager():
 
 def test_context_loads_values_from_config(monkeypatch):
     subsection = Config(password="1234")
-    config = Config(context=Config(subsection=subsection, my_key="my_value"))
+    config = Config(context=Config(subsection=subsection, key1="val1", key2="val2"))
     monkeypatch.setattr(prefect.utilities.context, "config", config)
-    fresh_context = Context()
+
+    fresh_context = Context(key2="new")
     assert "subsection" in fresh_context
-    assert fresh_context.my_key == "my_value"
+    assert fresh_context.key1 == "val1"
+    assert fresh_context.key2 == "new"  # overridden by constructor
     assert fresh_context.subsection == subsection
 
 
@@ -137,6 +139,27 @@ def test_context_loads_secrets_from_config(monkeypatch):
     assert fresh_context.secrets == secrets_dict
 
 
+def test_context_contextmanager_prioritizes_new_keys_even_on_context_exit(monkeypatch):
+    """Previously exiting a context block would reload from the config,
+    overwriting any explicitly set values in a nested context. This was due to
+    the `Context` constructor being implicitly called when stashing the old
+    context, and the constructor prioritizing `config.context` over explicit
+    values."""
+    config = Config(context=Config(my_key="fizz"))
+    monkeypatch.setattr(prefect.utilities.context, "config", config)
+
+    context = Context()
+    assert context.my_key == "fizz"
+
+    with context(my_key="buzz"):
+        assert context.my_key == "buzz"
+        with context({"config": {"logging": {"log_to_cloud": "FOO"}}}):
+            assert context.config.logging.log_to_cloud == "FOO"
+            assert context.my_key == "buzz"
+        assert context.my_key == "buzz"
+    assert context.my_key == "fizz"
+
+
 def test_context_contextmanager_prioritizes_new_config_keys():
     with prefect.context({"config": {"logging": {"log_to_cloud": "FOO"}}}):
         assert prefect.context.config.logging.log_to_cloud == "FOO"